SupritiVijay commited on
Commit
77d00df
·
verified ·
1 Parent(s): 4ac1eb7

Upload folder using huggingface_hub

Browse files
This view is limited to 50 files because it contains too many changes.   See raw diff
Files changed (50) hide show
  1. openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1025/config.json +39 -0
  2. openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1025/generation_config.json +12 -0
  3. openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1025/latest +1 -0
  4. openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1025/model.safetensors +3 -0
  5. openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1025/scheduler.pt +3 -0
  6. openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1025/trainer_state.json +748 -0
  7. openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1025/training_args.bin +3 -0
  8. openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1025/zero_to_fp32.py +760 -0
  9. openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1230/config.json +39 -0
  10. openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1230/generation_config.json +12 -0
  11. openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1230/latest +1 -0
  12. openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1230/model.safetensors +3 -0
  13. openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1230/scheduler.pt +3 -0
  14. openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1230/trainer_state.json +895 -0
  15. openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1230/training_args.bin +3 -0
  16. openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1230/zero_to_fp32.py +760 -0
  17. openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1435/config.json +39 -0
  18. openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1435/generation_config.json +12 -0
  19. openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1435/latest +1 -0
  20. openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1435/model.safetensors +3 -0
  21. openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1435/scheduler.pt +3 -0
  22. openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1435/trainer_state.json +1035 -0
  23. openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1435/training_args.bin +3 -0
  24. openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1435/zero_to_fp32.py +760 -0
  25. openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1640/config.json +39 -0
  26. openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1640/generation_config.json +12 -0
  27. openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1640/latest +1 -0
  28. openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1640/model.safetensors +3 -0
  29. openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1640/scheduler.pt +3 -0
  30. openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1640/trainer_state.json +1182 -0
  31. openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1640/training_args.bin +3 -0
  32. openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1640/zero_to_fp32.py +760 -0
  33. openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1845/config.json +39 -0
  34. openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1845/generation_config.json +12 -0
  35. openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1845/latest +1 -0
  36. openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1845/model.safetensors +3 -0
  37. openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1845/scheduler.pt +3 -0
  38. openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1845/trainer_state.json +1322 -0
  39. openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1845/training_args.bin +3 -0
  40. openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1845/zero_to_fp32.py +760 -0
  41. openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-205/config.json +39 -0
  42. openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-205/generation_config.json +12 -0
  43. openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-205/latest +1 -0
  44. openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-205/model.safetensors +3 -0
  45. openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-205/scheduler.pt +3 -0
  46. openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-205/trainer_state.json +174 -0
  47. openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-205/training_args.bin +3 -0
  48. openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-205/zero_to_fp32.py +760 -0
  49. openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-2050/config.json +39 -0
  50. openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-2050/generation_config.json +12 -0
openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1025/config.json ADDED
@@ -0,0 +1,39 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "architectures": [
3
+ "LlamaForCausalLM"
4
+ ],
5
+ "attention_bias": false,
6
+ "attention_dropout": 0.0,
7
+ "bos_token_id": 128000,
8
+ "dtype": "bfloat16",
9
+ "eos_token_id": [
10
+ 128001,
11
+ 128008,
12
+ 128009
13
+ ],
14
+ "head_dim": 64,
15
+ "hidden_act": "silu",
16
+ "hidden_size": 2048,
17
+ "initializer_range": 0.02,
18
+ "intermediate_size": 8192,
19
+ "max_position_embeddings": 131072,
20
+ "mlp_bias": false,
21
+ "model_type": "llama",
22
+ "num_attention_heads": 32,
23
+ "num_hidden_layers": 16,
24
+ "num_key_value_heads": 8,
25
+ "pretraining_tp": 1,
26
+ "rms_norm_eps": 1e-05,
27
+ "rope_scaling": {
28
+ "factor": 32.0,
29
+ "high_freq_factor": 4.0,
30
+ "low_freq_factor": 1.0,
31
+ "original_max_position_embeddings": 8192,
32
+ "rope_type": "llama3"
33
+ },
34
+ "rope_theta": 500000.0,
35
+ "tie_word_embeddings": true,
36
+ "transformers_version": "4.57.1",
37
+ "use_cache": false,
38
+ "vocab_size": 128256
39
+ }
openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1025/generation_config.json ADDED
@@ -0,0 +1,12 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "bos_token_id": 128000,
3
+ "do_sample": true,
4
+ "eos_token_id": [
5
+ 128001,
6
+ 128008,
7
+ 128009
8
+ ],
9
+ "temperature": 0.6,
10
+ "top_p": 0.9,
11
+ "transformers_version": "4.57.1"
12
+ }
openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1025/latest ADDED
@@ -0,0 +1 @@
 
 
1
+ global_step1025
openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1025/model.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:8a8893a0d101e747951d9f57c04f3fb50d463eb951ca1afaf23685df2bb15205
3
+ size 2996982344
openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1025/scheduler.pt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:3bfc7021ac48ba0d3d01a50fe2ad034e0b238728f73fd320e110d6f90de4b81a
3
+ size 1465
openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1025/trainer_state.json ADDED
@@ -0,0 +1,748 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "best_global_step": null,
3
+ "best_metric": null,
4
+ "best_model_checkpoint": null,
5
+ "epoch": 1.2486288848263254,
6
+ "eval_steps": 500,
7
+ "global_step": 1025,
8
+ "is_hyper_param_search": false,
9
+ "is_local_process_zero": true,
10
+ "is_world_process_zero": true,
11
+ "log_history": [
12
+ {
13
+ "epoch": 0.01218769043266301,
14
+ "grad_norm": 0.36486390233039856,
15
+ "learning_rate": 5e-05,
16
+ "loss": 1.2213,
17
+ "step": 10
18
+ },
19
+ {
20
+ "epoch": 0.02437538086532602,
21
+ "grad_norm": 0.2953564524650574,
22
+ "learning_rate": 5e-05,
23
+ "loss": 1.0223,
24
+ "step": 20
25
+ },
26
+ {
27
+ "epoch": 0.03656307129798903,
28
+ "grad_norm": 0.2766735851764679,
29
+ "learning_rate": 5e-05,
30
+ "loss": 0.9589,
31
+ "step": 30
32
+ },
33
+ {
34
+ "epoch": 0.04875076173065204,
35
+ "grad_norm": 0.32770997285842896,
36
+ "learning_rate": 5e-05,
37
+ "loss": 0.9081,
38
+ "step": 40
39
+ },
40
+ {
41
+ "epoch": 0.06093845216331505,
42
+ "grad_norm": 0.3052276372909546,
43
+ "learning_rate": 5e-05,
44
+ "loss": 0.8752,
45
+ "step": 50
46
+ },
47
+ {
48
+ "epoch": 0.07312614259597806,
49
+ "grad_norm": 0.240146666765213,
50
+ "learning_rate": 5e-05,
51
+ "loss": 0.873,
52
+ "step": 60
53
+ },
54
+ {
55
+ "epoch": 0.08531383302864107,
56
+ "grad_norm": 0.23872187733650208,
57
+ "learning_rate": 5e-05,
58
+ "loss": 0.8587,
59
+ "step": 70
60
+ },
61
+ {
62
+ "epoch": 0.09750152346130408,
63
+ "grad_norm": 0.32638487219810486,
64
+ "learning_rate": 5e-05,
65
+ "loss": 0.8656,
66
+ "step": 80
67
+ },
68
+ {
69
+ "epoch": 0.10968921389396709,
70
+ "grad_norm": 0.25019028782844543,
71
+ "learning_rate": 5e-05,
72
+ "loss": 0.8685,
73
+ "step": 90
74
+ },
75
+ {
76
+ "epoch": 0.1218769043266301,
77
+ "grad_norm": 0.2507675886154175,
78
+ "learning_rate": 5e-05,
79
+ "loss": 0.8662,
80
+ "step": 100
81
+ },
82
+ {
83
+ "epoch": 0.1340645947592931,
84
+ "grad_norm": 0.27966755628585815,
85
+ "learning_rate": 5e-05,
86
+ "loss": 0.8428,
87
+ "step": 110
88
+ },
89
+ {
90
+ "epoch": 0.14625228519195613,
91
+ "grad_norm": 0.2622673213481903,
92
+ "learning_rate": 5e-05,
93
+ "loss": 0.8416,
94
+ "step": 120
95
+ },
96
+ {
97
+ "epoch": 0.15843997562461914,
98
+ "grad_norm": 0.3003966510295868,
99
+ "learning_rate": 5e-05,
100
+ "loss": 0.8528,
101
+ "step": 130
102
+ },
103
+ {
104
+ "epoch": 0.17062766605728213,
105
+ "grad_norm": 0.32308319211006165,
106
+ "learning_rate": 5e-05,
107
+ "loss": 0.8357,
108
+ "step": 140
109
+ },
110
+ {
111
+ "epoch": 0.18281535648994515,
112
+ "grad_norm": 0.27279096841812134,
113
+ "learning_rate": 5e-05,
114
+ "loss": 0.8432,
115
+ "step": 150
116
+ },
117
+ {
118
+ "epoch": 0.19500304692260817,
119
+ "grad_norm": 0.27976104617118835,
120
+ "learning_rate": 5e-05,
121
+ "loss": 0.8119,
122
+ "step": 160
123
+ },
124
+ {
125
+ "epoch": 0.2071907373552712,
126
+ "grad_norm": 0.290866881608963,
127
+ "learning_rate": 5e-05,
128
+ "loss": 0.8096,
129
+ "step": 170
130
+ },
131
+ {
132
+ "epoch": 0.21937842778793418,
133
+ "grad_norm": 0.29430949687957764,
134
+ "learning_rate": 5e-05,
135
+ "loss": 0.819,
136
+ "step": 180
137
+ },
138
+ {
139
+ "epoch": 0.2315661182205972,
140
+ "grad_norm": 0.27714136242866516,
141
+ "learning_rate": 5e-05,
142
+ "loss": 0.8187,
143
+ "step": 190
144
+ },
145
+ {
146
+ "epoch": 0.2437538086532602,
147
+ "grad_norm": 0.2997029423713684,
148
+ "learning_rate": 5e-05,
149
+ "loss": 0.808,
150
+ "step": 200
151
+ },
152
+ {
153
+ "epoch": 0.25594149908592323,
154
+ "grad_norm": 0.2539573609828949,
155
+ "learning_rate": 5e-05,
156
+ "loss": 0.7943,
157
+ "step": 210
158
+ },
159
+ {
160
+ "epoch": 0.2681291895185862,
161
+ "grad_norm": 0.23294171690940857,
162
+ "learning_rate": 5e-05,
163
+ "loss": 0.8188,
164
+ "step": 220
165
+ },
166
+ {
167
+ "epoch": 0.28031687995124926,
168
+ "grad_norm": 0.25016292929649353,
169
+ "learning_rate": 5e-05,
170
+ "loss": 0.8035,
171
+ "step": 230
172
+ },
173
+ {
174
+ "epoch": 0.29250457038391225,
175
+ "grad_norm": 0.37125322222709656,
176
+ "learning_rate": 5e-05,
177
+ "loss": 0.785,
178
+ "step": 240
179
+ },
180
+ {
181
+ "epoch": 0.30469226081657524,
182
+ "grad_norm": 0.21520452201366425,
183
+ "learning_rate": 5e-05,
184
+ "loss": 0.8088,
185
+ "step": 250
186
+ },
187
+ {
188
+ "epoch": 0.3168799512492383,
189
+ "grad_norm": 0.20990602672100067,
190
+ "learning_rate": 5e-05,
191
+ "loss": 0.7813,
192
+ "step": 260
193
+ },
194
+ {
195
+ "epoch": 0.3290676416819013,
196
+ "grad_norm": 0.24193961918354034,
197
+ "learning_rate": 5e-05,
198
+ "loss": 0.7979,
199
+ "step": 270
200
+ },
201
+ {
202
+ "epoch": 0.34125533211456427,
203
+ "grad_norm": 0.22766050696372986,
204
+ "learning_rate": 5e-05,
205
+ "loss": 0.785,
206
+ "step": 280
207
+ },
208
+ {
209
+ "epoch": 0.3534430225472273,
210
+ "grad_norm": 0.2662702202796936,
211
+ "learning_rate": 5e-05,
212
+ "loss": 0.7766,
213
+ "step": 290
214
+ },
215
+ {
216
+ "epoch": 0.3656307129798903,
217
+ "grad_norm": 0.23020495474338531,
218
+ "learning_rate": 5e-05,
219
+ "loss": 0.7687,
220
+ "step": 300
221
+ },
222
+ {
223
+ "epoch": 0.3778184034125533,
224
+ "grad_norm": 0.21770520508289337,
225
+ "learning_rate": 5e-05,
226
+ "loss": 0.7815,
227
+ "step": 310
228
+ },
229
+ {
230
+ "epoch": 0.39000609384521634,
231
+ "grad_norm": 0.22677403688430786,
232
+ "learning_rate": 5e-05,
233
+ "loss": 0.7999,
234
+ "step": 320
235
+ },
236
+ {
237
+ "epoch": 0.40219378427787933,
238
+ "grad_norm": 0.24891860783100128,
239
+ "learning_rate": 5e-05,
240
+ "loss": 0.7946,
241
+ "step": 330
242
+ },
243
+ {
244
+ "epoch": 0.4143814747105424,
245
+ "grad_norm": 0.2506854832172394,
246
+ "learning_rate": 5e-05,
247
+ "loss": 0.7889,
248
+ "step": 340
249
+ },
250
+ {
251
+ "epoch": 0.42656916514320536,
252
+ "grad_norm": 0.24388180673122406,
253
+ "learning_rate": 5e-05,
254
+ "loss": 0.79,
255
+ "step": 350
256
+ },
257
+ {
258
+ "epoch": 0.43875685557586835,
259
+ "grad_norm": 0.22329650819301605,
260
+ "learning_rate": 5e-05,
261
+ "loss": 0.7632,
262
+ "step": 360
263
+ },
264
+ {
265
+ "epoch": 0.4509445460085314,
266
+ "grad_norm": 0.2240927815437317,
267
+ "learning_rate": 5e-05,
268
+ "loss": 0.8042,
269
+ "step": 370
270
+ },
271
+ {
272
+ "epoch": 0.4631322364411944,
273
+ "grad_norm": 0.2578203082084656,
274
+ "learning_rate": 5e-05,
275
+ "loss": 0.7667,
276
+ "step": 380
277
+ },
278
+ {
279
+ "epoch": 0.4753199268738574,
280
+ "grad_norm": 0.25349217653274536,
281
+ "learning_rate": 5e-05,
282
+ "loss": 0.7911,
283
+ "step": 390
284
+ },
285
+ {
286
+ "epoch": 0.4875076173065204,
287
+ "grad_norm": 0.254446804523468,
288
+ "learning_rate": 5e-05,
289
+ "loss": 0.7909,
290
+ "step": 400
291
+ },
292
+ {
293
+ "epoch": 0.4996953077391834,
294
+ "grad_norm": 0.24253171682357788,
295
+ "learning_rate": 5e-05,
296
+ "loss": 0.7896,
297
+ "step": 410
298
+ },
299
+ {
300
+ "epoch": 0.5118829981718465,
301
+ "grad_norm": 0.27901458740234375,
302
+ "learning_rate": 5e-05,
303
+ "loss": 0.7828,
304
+ "step": 420
305
+ },
306
+ {
307
+ "epoch": 0.5240706886045094,
308
+ "grad_norm": 0.31884434819221497,
309
+ "learning_rate": 5e-05,
310
+ "loss": 0.7919,
311
+ "step": 430
312
+ },
313
+ {
314
+ "epoch": 0.5362583790371724,
315
+ "grad_norm": 0.22920799255371094,
316
+ "learning_rate": 5e-05,
317
+ "loss": 0.7777,
318
+ "step": 440
319
+ },
320
+ {
321
+ "epoch": 0.5484460694698354,
322
+ "grad_norm": 0.22606618702411652,
323
+ "learning_rate": 5e-05,
324
+ "loss": 0.7661,
325
+ "step": 450
326
+ },
327
+ {
328
+ "epoch": 0.5606337599024985,
329
+ "grad_norm": 0.21604052186012268,
330
+ "learning_rate": 5e-05,
331
+ "loss": 0.769,
332
+ "step": 460
333
+ },
334
+ {
335
+ "epoch": 0.5728214503351615,
336
+ "grad_norm": 0.2296101152896881,
337
+ "learning_rate": 5e-05,
338
+ "loss": 0.7626,
339
+ "step": 470
340
+ },
341
+ {
342
+ "epoch": 0.5850091407678245,
343
+ "grad_norm": 0.24162150919437408,
344
+ "learning_rate": 5e-05,
345
+ "loss": 0.7831,
346
+ "step": 480
347
+ },
348
+ {
349
+ "epoch": 0.5971968312004875,
350
+ "grad_norm": 0.22364799678325653,
351
+ "learning_rate": 5e-05,
352
+ "loss": 0.7737,
353
+ "step": 490
354
+ },
355
+ {
356
+ "epoch": 0.6093845216331505,
357
+ "grad_norm": 0.23953403532505035,
358
+ "learning_rate": 5e-05,
359
+ "loss": 0.787,
360
+ "step": 500
361
+ },
362
+ {
363
+ "epoch": 0.6215722120658135,
364
+ "grad_norm": 0.22349311411380768,
365
+ "learning_rate": 5e-05,
366
+ "loss": 0.7763,
367
+ "step": 510
368
+ },
369
+ {
370
+ "epoch": 0.6337599024984766,
371
+ "grad_norm": 0.21996377408504486,
372
+ "learning_rate": 5e-05,
373
+ "loss": 0.7596,
374
+ "step": 520
375
+ },
376
+ {
377
+ "epoch": 0.6459475929311396,
378
+ "grad_norm": 0.245538130402565,
379
+ "learning_rate": 5e-05,
380
+ "loss": 0.7785,
381
+ "step": 530
382
+ },
383
+ {
384
+ "epoch": 0.6581352833638026,
385
+ "grad_norm": 0.23442520201206207,
386
+ "learning_rate": 5e-05,
387
+ "loss": 0.7633,
388
+ "step": 540
389
+ },
390
+ {
391
+ "epoch": 0.6703229737964655,
392
+ "grad_norm": 0.22274106740951538,
393
+ "learning_rate": 5e-05,
394
+ "loss": 0.7756,
395
+ "step": 550
396
+ },
397
+ {
398
+ "epoch": 0.6825106642291285,
399
+ "grad_norm": 0.21717530488967896,
400
+ "learning_rate": 5e-05,
401
+ "loss": 0.7726,
402
+ "step": 560
403
+ },
404
+ {
405
+ "epoch": 0.6946983546617916,
406
+ "grad_norm": 0.2465955913066864,
407
+ "learning_rate": 5e-05,
408
+ "loss": 0.7476,
409
+ "step": 570
410
+ },
411
+ {
412
+ "epoch": 0.7068860450944546,
413
+ "grad_norm": 0.2318575084209442,
414
+ "learning_rate": 5e-05,
415
+ "loss": 0.7584,
416
+ "step": 580
417
+ },
418
+ {
419
+ "epoch": 0.7190737355271176,
420
+ "grad_norm": 0.25114378333091736,
421
+ "learning_rate": 5e-05,
422
+ "loss": 0.7625,
423
+ "step": 590
424
+ },
425
+ {
426
+ "epoch": 0.7312614259597806,
427
+ "grad_norm": 0.19145755469799042,
428
+ "learning_rate": 5e-05,
429
+ "loss": 0.7577,
430
+ "step": 600
431
+ },
432
+ {
433
+ "epoch": 0.7434491163924436,
434
+ "grad_norm": 0.19129066169261932,
435
+ "learning_rate": 5e-05,
436
+ "loss": 0.7772,
437
+ "step": 610
438
+ },
439
+ {
440
+ "epoch": 0.7556368068251066,
441
+ "grad_norm": 0.2014475166797638,
442
+ "learning_rate": 5e-05,
443
+ "loss": 0.7601,
444
+ "step": 620
445
+ },
446
+ {
447
+ "epoch": 0.7678244972577697,
448
+ "grad_norm": 0.21480509638786316,
449
+ "learning_rate": 5e-05,
450
+ "loss": 0.7665,
451
+ "step": 630
452
+ },
453
+ {
454
+ "epoch": 0.7800121876904327,
455
+ "grad_norm": 0.2633228302001953,
456
+ "learning_rate": 5e-05,
457
+ "loss": 0.7565,
458
+ "step": 640
459
+ },
460
+ {
461
+ "epoch": 0.7921998781230957,
462
+ "grad_norm": 0.24633203446865082,
463
+ "learning_rate": 5e-05,
464
+ "loss": 0.7602,
465
+ "step": 650
466
+ },
467
+ {
468
+ "epoch": 0.8043875685557587,
469
+ "grad_norm": 0.19770534336566925,
470
+ "learning_rate": 5e-05,
471
+ "loss": 0.7685,
472
+ "step": 660
473
+ },
474
+ {
475
+ "epoch": 0.8165752589884216,
476
+ "grad_norm": 0.2267930805683136,
477
+ "learning_rate": 5e-05,
478
+ "loss": 0.7516,
479
+ "step": 670
480
+ },
481
+ {
482
+ "epoch": 0.8287629494210847,
483
+ "grad_norm": 0.238111674785614,
484
+ "learning_rate": 5e-05,
485
+ "loss": 0.7487,
486
+ "step": 680
487
+ },
488
+ {
489
+ "epoch": 0.8409506398537477,
490
+ "grad_norm": 0.19528904557228088,
491
+ "learning_rate": 5e-05,
492
+ "loss": 0.7467,
493
+ "step": 690
494
+ },
495
+ {
496
+ "epoch": 0.8531383302864107,
497
+ "grad_norm": 0.21305914223194122,
498
+ "learning_rate": 5e-05,
499
+ "loss": 0.7527,
500
+ "step": 700
501
+ },
502
+ {
503
+ "epoch": 0.8653260207190737,
504
+ "grad_norm": 0.2205205112695694,
505
+ "learning_rate": 5e-05,
506
+ "loss": 0.7639,
507
+ "step": 710
508
+ },
509
+ {
510
+ "epoch": 0.8775137111517367,
511
+ "grad_norm": 0.22429127991199493,
512
+ "learning_rate": 5e-05,
513
+ "loss": 0.7744,
514
+ "step": 720
515
+ },
516
+ {
517
+ "epoch": 0.8897014015843998,
518
+ "grad_norm": 0.23755772411823273,
519
+ "learning_rate": 5e-05,
520
+ "loss": 0.7446,
521
+ "step": 730
522
+ },
523
+ {
524
+ "epoch": 0.9018890920170628,
525
+ "grad_norm": 0.17592693865299225,
526
+ "learning_rate": 5e-05,
527
+ "loss": 0.7442,
528
+ "step": 740
529
+ },
530
+ {
531
+ "epoch": 0.9140767824497258,
532
+ "grad_norm": 0.2266898900270462,
533
+ "learning_rate": 5e-05,
534
+ "loss": 0.7326,
535
+ "step": 750
536
+ },
537
+ {
538
+ "epoch": 0.9262644728823888,
539
+ "grad_norm": 0.23541569709777832,
540
+ "learning_rate": 5e-05,
541
+ "loss": 0.7338,
542
+ "step": 760
543
+ },
544
+ {
545
+ "epoch": 0.9384521633150518,
546
+ "grad_norm": 0.22233888506889343,
547
+ "learning_rate": 5e-05,
548
+ "loss": 0.7385,
549
+ "step": 770
550
+ },
551
+ {
552
+ "epoch": 0.9506398537477148,
553
+ "grad_norm": 0.23714767396450043,
554
+ "learning_rate": 5e-05,
555
+ "loss": 0.7306,
556
+ "step": 780
557
+ },
558
+ {
559
+ "epoch": 0.9628275441803779,
560
+ "grad_norm": 0.21191060543060303,
561
+ "learning_rate": 5e-05,
562
+ "loss": 0.736,
563
+ "step": 790
564
+ },
565
+ {
566
+ "epoch": 0.9750152346130408,
567
+ "grad_norm": 0.22759106755256653,
568
+ "learning_rate": 5e-05,
569
+ "loss": 0.7505,
570
+ "step": 800
571
+ },
572
+ {
573
+ "epoch": 0.9872029250457038,
574
+ "grad_norm": 0.19895903766155243,
575
+ "learning_rate": 5e-05,
576
+ "loss": 0.7536,
577
+ "step": 810
578
+ },
579
+ {
580
+ "epoch": 0.9993906154783668,
581
+ "grad_norm": 0.3061685860157013,
582
+ "learning_rate": 5e-05,
583
+ "loss": 0.7493,
584
+ "step": 820
585
+ },
586
+ {
587
+ "epoch": 1.0109689213893966,
588
+ "grad_norm": 0.18868015706539154,
589
+ "learning_rate": 5e-05,
590
+ "loss": 0.7342,
591
+ "step": 830
592
+ },
593
+ {
594
+ "epoch": 1.0231566118220596,
595
+ "grad_norm": 0.1797628402709961,
596
+ "learning_rate": 5e-05,
597
+ "loss": 0.7192,
598
+ "step": 840
599
+ },
600
+ {
601
+ "epoch": 1.0353443022547226,
602
+ "grad_norm": 0.19511596858501434,
603
+ "learning_rate": 5e-05,
604
+ "loss": 0.7002,
605
+ "step": 850
606
+ },
607
+ {
608
+ "epoch": 1.0475319926873858,
609
+ "grad_norm": 0.2123185396194458,
610
+ "learning_rate": 5e-05,
611
+ "loss": 0.7236,
612
+ "step": 860
613
+ },
614
+ {
615
+ "epoch": 1.0597196831200488,
616
+ "grad_norm": 0.2047470211982727,
617
+ "learning_rate": 5e-05,
618
+ "loss": 0.711,
619
+ "step": 870
620
+ },
621
+ {
622
+ "epoch": 1.0719073735527118,
623
+ "grad_norm": 0.18809868395328522,
624
+ "learning_rate": 5e-05,
625
+ "loss": 0.7253,
626
+ "step": 880
627
+ },
628
+ {
629
+ "epoch": 1.0840950639853748,
630
+ "grad_norm": 0.2207012176513672,
631
+ "learning_rate": 5e-05,
632
+ "loss": 0.6971,
633
+ "step": 890
634
+ },
635
+ {
636
+ "epoch": 1.0962827544180378,
637
+ "grad_norm": 0.1912035495042801,
638
+ "learning_rate": 5e-05,
639
+ "loss": 0.7239,
640
+ "step": 900
641
+ },
642
+ {
643
+ "epoch": 1.1084704448507008,
644
+ "grad_norm": 0.23486360907554626,
645
+ "learning_rate": 5e-05,
646
+ "loss": 0.7335,
647
+ "step": 910
648
+ },
649
+ {
650
+ "epoch": 1.1206581352833638,
651
+ "grad_norm": 0.22189787030220032,
652
+ "learning_rate": 5e-05,
653
+ "loss": 0.7149,
654
+ "step": 920
655
+ },
656
+ {
657
+ "epoch": 1.1328458257160268,
658
+ "grad_norm": 0.25818824768066406,
659
+ "learning_rate": 5e-05,
660
+ "loss": 0.7322,
661
+ "step": 930
662
+ },
663
+ {
664
+ "epoch": 1.1450335161486898,
665
+ "grad_norm": 0.2088775783777237,
666
+ "learning_rate": 5e-05,
667
+ "loss": 0.7382,
668
+ "step": 940
669
+ },
670
+ {
671
+ "epoch": 1.1572212065813527,
672
+ "grad_norm": 0.19445833563804626,
673
+ "learning_rate": 5e-05,
674
+ "loss": 0.7223,
675
+ "step": 950
676
+ },
677
+ {
678
+ "epoch": 1.1694088970140157,
679
+ "grad_norm": 0.21489456295967102,
680
+ "learning_rate": 5e-05,
681
+ "loss": 0.7317,
682
+ "step": 960
683
+ },
684
+ {
685
+ "epoch": 1.181596587446679,
686
+ "grad_norm": 0.24050794541835785,
687
+ "learning_rate": 5e-05,
688
+ "loss": 0.7275,
689
+ "step": 970
690
+ },
691
+ {
692
+ "epoch": 1.193784277879342,
693
+ "grad_norm": 0.18142102658748627,
694
+ "learning_rate": 5e-05,
695
+ "loss": 0.7046,
696
+ "step": 980
697
+ },
698
+ {
699
+ "epoch": 1.205971968312005,
700
+ "grad_norm": 0.4781077802181244,
701
+ "learning_rate": 5e-05,
702
+ "loss": 0.7433,
703
+ "step": 990
704
+ },
705
+ {
706
+ "epoch": 1.218159658744668,
707
+ "grad_norm": 0.20628023147583008,
708
+ "learning_rate": 5e-05,
709
+ "loss": 0.7403,
710
+ "step": 1000
711
+ },
712
+ {
713
+ "epoch": 1.230347349177331,
714
+ "grad_norm": 0.19117140769958496,
715
+ "learning_rate": 5e-05,
716
+ "loss": 0.7308,
717
+ "step": 1010
718
+ },
719
+ {
720
+ "epoch": 1.242535039609994,
721
+ "grad_norm": 0.2107745110988617,
722
+ "learning_rate": 5e-05,
723
+ "loss": 0.7201,
724
+ "step": 1020
725
+ }
726
+ ],
727
+ "logging_steps": 10,
728
+ "max_steps": 2463,
729
+ "num_input_tokens_seen": 0,
730
+ "num_train_epochs": 3,
731
+ "save_steps": 205,
732
+ "stateful_callbacks": {
733
+ "TrainerControl": {
734
+ "args": {
735
+ "should_epoch_stop": false,
736
+ "should_evaluate": false,
737
+ "should_log": false,
738
+ "should_save": true,
739
+ "should_training_stop": false
740
+ },
741
+ "attributes": {}
742
+ }
743
+ },
744
+ "total_flos": 4.207123313442423e+18,
745
+ "train_batch_size": 4,
746
+ "trial_name": null,
747
+ "trial_params": null
748
+ }
openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1025/training_args.bin ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:56896ea21940e33d7e2934267dff2a5b4a4a80195a0f165245f6617a0f88b90b
3
+ size 7377
openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1025/zero_to_fp32.py ADDED
@@ -0,0 +1,760 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python
2
+
3
+ # Copyright (c) Microsoft Corporation.
4
+ # SPDX-License-Identifier: Apache-2.0
5
+
6
+ # DeepSpeed Team
7
+
8
+ # This script extracts fp32 consolidated weights from a zero 1, 2 and 3 DeepSpeed checkpoints. It gets
9
+ # copied into the top level checkpoint dir, so the user can easily do the conversion at any point in
10
+ # the future. Once extracted, the weights don't require DeepSpeed and can be used in any
11
+ # application.
12
+ #
13
+ # example:
14
+ # python zero_to_fp32.py . output_dir/
15
+ # or
16
+ # python zero_to_fp32.py . output_dir/ --safe_serialization
17
+
18
+ import argparse
19
+ import torch
20
+ import glob
21
+ import math
22
+ import os
23
+ import re
24
+ import gc
25
+ import json
26
+ import numpy as np
27
+ from tqdm import tqdm
28
+ from collections import OrderedDict
29
+ from dataclasses import dataclass
30
+
31
+ # while this script doesn't use deepspeed to recover data, since the checkpoints are pickled with
32
+ # DeepSpeed data structures it has to be available in the current python environment.
33
+ from deepspeed.utils import logger
34
+ from deepspeed.checkpoint.constants import (DS_VERSION, OPTIMIZER_STATE_DICT, SINGLE_PARTITION_OF_FP32_GROUPS,
35
+ FP32_FLAT_GROUPS, ZERO_STAGE, PARTITION_COUNT, PARAM_SHAPES, BUFFER_NAMES,
36
+ FROZEN_PARAM_SHAPES, FROZEN_PARAM_FRAGMENTS)
37
+
38
+
39
+ @dataclass
40
+ class zero_model_state:
41
+ buffers: dict()
42
+ param_shapes: dict()
43
+ shared_params: list
44
+ ds_version: int
45
+ frozen_param_shapes: dict()
46
+ frozen_param_fragments: dict()
47
+
48
+
49
+ debug = 0
50
+
51
+ # load to cpu
52
+ device = torch.device('cpu')
53
+
54
+
55
+ def atoi(text):
56
+ return int(text) if text.isdigit() else text
57
+
58
+
59
+ def natural_keys(text):
60
+ '''
61
+ alist.sort(key=natural_keys) sorts in human order
62
+ http://nedbatchelder.com/blog/200712/human_sorting.html
63
+ (See Toothy's implementation in the comments)
64
+ '''
65
+ return [atoi(c) for c in re.split(r'(\d+)', text)]
66
+
67
+
68
+ def get_model_state_file(checkpoint_dir, zero_stage):
69
+ if not os.path.isdir(checkpoint_dir):
70
+ raise FileNotFoundError(f"Directory '{checkpoint_dir}' doesn't exist")
71
+
72
+ # there should be only one file
73
+ if zero_stage <= 2:
74
+ file = os.path.join(checkpoint_dir, "mp_rank_00_model_states.pt")
75
+ elif zero_stage == 3:
76
+ file = os.path.join(checkpoint_dir, "zero_pp_rank_0_mp_rank_00_model_states.pt")
77
+
78
+ if not os.path.exists(file):
79
+ raise FileNotFoundError(f"can't find model states file at '{file}'")
80
+
81
+ return file
82
+
83
+
84
+ def get_checkpoint_files(checkpoint_dir, glob_pattern):
85
+ # XXX: need to test that this simple glob rule works for multi-node setup too
86
+ ckpt_files = sorted(glob.glob(os.path.join(checkpoint_dir, glob_pattern)), key=natural_keys)
87
+
88
+ if len(ckpt_files) == 0:
89
+ raise FileNotFoundError(f"can't find {glob_pattern} files in directory '{checkpoint_dir}'")
90
+
91
+ return ckpt_files
92
+
93
+
94
+ def get_optim_files(checkpoint_dir):
95
+ return get_checkpoint_files(checkpoint_dir, "*_optim_states.pt")
96
+
97
+
98
+ def get_model_state_files(checkpoint_dir):
99
+ return get_checkpoint_files(checkpoint_dir, "*_model_states.pt")
100
+
101
+
102
+ def parse_model_states(files):
103
+ zero_model_states = []
104
+ for file in files:
105
+ state_dict = torch.load(file, map_location=device, weights_only=False)
106
+
107
+ if BUFFER_NAMES not in state_dict:
108
+ raise ValueError(f"{file} is not a model state checkpoint")
109
+ buffer_names = state_dict[BUFFER_NAMES]
110
+ if debug:
111
+ print("Found buffers:", buffer_names)
112
+
113
+ # recover just the buffers while restoring them to fp32 if they were saved in fp16
114
+ buffers = {k: v.float() for k, v in state_dict["module"].items() if k in buffer_names}
115
+ param_shapes = state_dict[PARAM_SHAPES]
116
+
117
+ # collect parameters that are included in param_shapes
118
+ param_names = []
119
+ for s in param_shapes:
120
+ for name in s.keys():
121
+ param_names.append(name)
122
+
123
+ # update with frozen parameters
124
+ frozen_param_shapes = state_dict.get(FROZEN_PARAM_SHAPES, None)
125
+ if frozen_param_shapes is not None:
126
+ if debug:
127
+ print(f"Found frozen_param_shapes: {frozen_param_shapes}")
128
+ param_names += list(frozen_param_shapes.keys())
129
+
130
+ # handle shared params
131
+ shared_params = [[k, v] for k, v in state_dict["shared_params"].items()]
132
+
133
+ ds_version = state_dict.get(DS_VERSION, None)
134
+
135
+ frozen_param_fragments = state_dict.get(FROZEN_PARAM_FRAGMENTS, None)
136
+
137
+ z_model_state = zero_model_state(buffers=buffers,
138
+ param_shapes=param_shapes,
139
+ shared_params=shared_params,
140
+ ds_version=ds_version,
141
+ frozen_param_shapes=frozen_param_shapes,
142
+ frozen_param_fragments=frozen_param_fragments)
143
+ zero_model_states.append(z_model_state)
144
+
145
+ return zero_model_states
146
+
147
+
148
+ def parse_optim_states(files, ds_checkpoint_dir):
149
+ total_files = len(files)
150
+ state_dicts = []
151
+ for f in tqdm(files, desc='Loading checkpoint shards'):
152
+ state_dict = torch.load(f, map_location=device, mmap=True, weights_only=False)
153
+ # immediately discard the potentially huge 2 optimizer states as we only care for fp32 master weights
154
+ # and also handle the case where it was already removed by another helper script
155
+ state_dict["optimizer_state_dict"].pop("optimizer_state_dict", None)
156
+ state_dicts.append(state_dict)
157
+
158
+ if ZERO_STAGE not in state_dicts[0][OPTIMIZER_STATE_DICT]:
159
+ raise ValueError(f"{files[0]} is not a zero checkpoint")
160
+ zero_stage = state_dicts[0][OPTIMIZER_STATE_DICT][ZERO_STAGE]
161
+ world_size = state_dicts[0][OPTIMIZER_STATE_DICT][PARTITION_COUNT]
162
+
163
+ # For ZeRO-2 each param group can have different partition_count as data parallelism for expert
164
+ # parameters can be different from data parallelism for non-expert parameters. So we can just
165
+ # use the max of the partition_count to get the dp world_size.
166
+
167
+ if type(world_size) is list:
168
+ world_size = max(world_size)
169
+
170
+ if world_size != total_files:
171
+ raise ValueError(
172
+ f"Expected {world_size} of '*_optim_states.pt' under '{ds_checkpoint_dir}' but found {total_files} files. "
173
+ "Possibly due to an overwrite of an old checkpoint, or a checkpoint didn't get saved by one or more processes."
174
+ )
175
+
176
+ # the groups are named differently in each stage
177
+ if zero_stage <= 2:
178
+ fp32_groups_key = SINGLE_PARTITION_OF_FP32_GROUPS
179
+ elif zero_stage == 3:
180
+ fp32_groups_key = FP32_FLAT_GROUPS
181
+ else:
182
+ raise ValueError(f"unknown zero stage {zero_stage}")
183
+
184
+ fp32_flat_groups = [state_dicts[i][OPTIMIZER_STATE_DICT][fp32_groups_key] for i in range(len(state_dicts))]
185
+ return zero_stage, world_size, fp32_flat_groups
186
+
187
+
188
+ def _get_fp32_state_dict_from_zero_checkpoint(ds_checkpoint_dir, exclude_frozen_parameters):
189
+ """
190
+ Returns fp32 state_dict reconstructed from ds checkpoint
191
+
192
+ Args:
193
+ - ``ds_checkpoint_dir``: path to the deepspeed checkpoint folder (where the optimizer files are)
194
+
195
+ """
196
+ print(f"Processing zero checkpoint '{ds_checkpoint_dir}'")
197
+
198
+ optim_files = get_optim_files(ds_checkpoint_dir)
199
+ zero_stage, world_size, fp32_flat_groups = parse_optim_states(optim_files, ds_checkpoint_dir)
200
+ print(f"Detected checkpoint of type zero stage {zero_stage}, world_size: {world_size}")
201
+
202
+ model_files = get_model_state_files(ds_checkpoint_dir)
203
+
204
+ zero_model_states = parse_model_states(model_files)
205
+ print(f'Parsing checkpoint created by deepspeed=={zero_model_states[0].ds_version}')
206
+
207
+ if zero_stage <= 2:
208
+ return _get_fp32_state_dict_from_zero2_checkpoint(world_size, fp32_flat_groups, zero_model_states,
209
+ exclude_frozen_parameters)
210
+ elif zero_stage == 3:
211
+ return _get_fp32_state_dict_from_zero3_checkpoint(world_size, fp32_flat_groups, zero_model_states,
212
+ exclude_frozen_parameters)
213
+
214
+
215
+ def _zero2_merge_frozen_params(state_dict, zero_model_states):
216
+ if zero_model_states[0].frozen_param_shapes is None or len(zero_model_states[0].frozen_param_shapes) == 0:
217
+ return
218
+
219
+ frozen_param_shapes = zero_model_states[0].frozen_param_shapes
220
+ frozen_param_fragments = zero_model_states[0].frozen_param_fragments
221
+
222
+ if debug:
223
+ num_elem = sum(s.numel() for s in frozen_param_shapes.values())
224
+ print(f'rank 0: {FROZEN_PARAM_SHAPES}.numel = {num_elem}')
225
+
226
+ wanted_params = len(frozen_param_shapes)
227
+ wanted_numel = sum(s.numel() for s in frozen_param_shapes.values())
228
+ avail_numel = sum([p.numel() for p in frozen_param_fragments.values()])
229
+ print(f'Frozen params: Have {avail_numel} numels to process.')
230
+ print(f'Frozen params: Need {wanted_numel} numels in {wanted_params} params')
231
+
232
+ total_params = 0
233
+ total_numel = 0
234
+ for name, shape in frozen_param_shapes.items():
235
+ total_params += 1
236
+ unpartitioned_numel = shape.numel()
237
+ total_numel += unpartitioned_numel
238
+
239
+ state_dict[name] = frozen_param_fragments[name]
240
+
241
+ if debug:
242
+ print(f"{name} full shape: {shape} unpartitioned numel {unpartitioned_numel} ")
243
+
244
+ print(f"Reconstructed Frozen fp32 state dict with {total_params} params {total_numel} elements")
245
+
246
+
247
+ def _has_callable(obj, fn):
248
+ attr = getattr(obj, fn, None)
249
+ return callable(attr)
250
+
251
+
252
+ def _zero2_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states):
253
+ param_shapes = zero_model_states[0].param_shapes
254
+
255
+ # Reconstruction protocol:
256
+ #
257
+ # XXX: document this
258
+
259
+ if debug:
260
+ for i in range(world_size):
261
+ for j in range(len(fp32_flat_groups[0])):
262
+ print(f"{FP32_FLAT_GROUPS}[{i}][{j}].shape={fp32_flat_groups[i][j].shape}")
263
+
264
+ # XXX: memory usage doubles here (zero2)
265
+ num_param_groups = len(fp32_flat_groups[0])
266
+ merged_single_partition_of_fp32_groups = []
267
+ for i in range(num_param_groups):
268
+ merged_partitions = [sd[i] for sd in fp32_flat_groups]
269
+ full_single_fp32_vector = torch.cat(merged_partitions, 0)
270
+ merged_single_partition_of_fp32_groups.append(full_single_fp32_vector)
271
+ avail_numel = sum(
272
+ [full_single_fp32_vector.numel() for full_single_fp32_vector in merged_single_partition_of_fp32_groups])
273
+
274
+ if debug:
275
+ wanted_params = sum([len(shapes) for shapes in param_shapes])
276
+ wanted_numel = sum([sum(shape.numel() for shape in shapes.values()) for shapes in param_shapes])
277
+ # not asserting if there is a mismatch due to possible padding
278
+ print(f"Have {avail_numel} numels to process.")
279
+ print(f"Need {wanted_numel} numels in {wanted_params} params.")
280
+
281
+ # params
282
+ # XXX: for huge models that can't fit into the host's RAM we will have to recode this to support
283
+ # out-of-core computing solution
284
+ total_numel = 0
285
+ total_params = 0
286
+ for shapes, full_single_fp32_vector in zip(param_shapes, merged_single_partition_of_fp32_groups):
287
+ offset = 0
288
+ avail_numel = full_single_fp32_vector.numel()
289
+ for name, shape in shapes.items():
290
+
291
+ unpartitioned_numel = shape.numel() if _has_callable(shape, 'numel') else math.prod(shape)
292
+ total_numel += unpartitioned_numel
293
+ total_params += 1
294
+
295
+ if debug:
296
+ print(f"{name} full shape: {shape} unpartitioned numel {unpartitioned_numel} ")
297
+ state_dict[name] = full_single_fp32_vector.narrow(0, offset, unpartitioned_numel).view(shape)
298
+ offset += unpartitioned_numel
299
+
300
+ # Z2 started to align to 2*world_size to improve nccl performance. Therefore both offset and
301
+ # avail_numel can differ by anywhere between 0..2*world_size. Due to two unrelated complex
302
+ # paddings performed in the code it's almost impossible to predict the exact numbers w/o the
303
+ # live optimizer object, so we are checking that the numbers are within the right range
304
+ align_to = 2 * world_size
305
+
306
+ def zero2_align(x):
307
+ return align_to * math.ceil(x / align_to)
308
+
309
+ if debug:
310
+ print(f"original offset={offset}, avail_numel={avail_numel}")
311
+
312
+ offset = zero2_align(offset)
313
+ avail_numel = zero2_align(avail_numel)
314
+
315
+ if debug:
316
+ print(f"aligned offset={offset}, avail_numel={avail_numel}")
317
+
318
+ # Sanity check
319
+ if offset != avail_numel:
320
+ raise ValueError(f"consumed {offset} numels out of {avail_numel} - something is wrong")
321
+
322
+ print(f"Reconstructed fp32 state dict with {total_params} params {total_numel} elements")
323
+
324
+
325
+ def _get_fp32_state_dict_from_zero2_checkpoint(world_size, fp32_flat_groups, zero_model_states,
326
+ exclude_frozen_parameters):
327
+ state_dict = OrderedDict()
328
+
329
+ # buffers
330
+ buffers = zero_model_states[0].buffers
331
+ state_dict.update(buffers)
332
+ if debug:
333
+ print(f"added {len(buffers)} buffers")
334
+
335
+ if not exclude_frozen_parameters:
336
+ _zero2_merge_frozen_params(state_dict, zero_model_states)
337
+
338
+ _zero2_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states)
339
+
340
+ # recover shared parameters
341
+ for pair in zero_model_states[0].shared_params:
342
+ if pair[1] in state_dict:
343
+ state_dict[pair[0]] = state_dict[pair[1]]
344
+
345
+ return state_dict
346
+
347
+
348
+ def zero3_partitioned_param_info(unpartitioned_numel, world_size):
349
+ remainder = unpartitioned_numel % world_size
350
+ padding_numel = (world_size - remainder) if remainder else 0
351
+ partitioned_numel = math.ceil(unpartitioned_numel / world_size)
352
+ return partitioned_numel, padding_numel
353
+
354
+
355
+ def _zero3_merge_frozen_params(state_dict, world_size, zero_model_states):
356
+ if zero_model_states[0].frozen_param_shapes is None or len(zero_model_states[0].frozen_param_shapes) == 0:
357
+ return
358
+
359
+ if debug:
360
+ for i in range(world_size):
361
+ num_elem = sum(s.numel() for s in zero_model_states[i].frozen_param_fragments.values())
362
+ print(f'rank {i}: {FROZEN_PARAM_SHAPES}.numel = {num_elem}')
363
+
364
+ frozen_param_shapes = zero_model_states[0].frozen_param_shapes
365
+ wanted_params = len(frozen_param_shapes)
366
+ wanted_numel = sum(s.numel() for s in frozen_param_shapes.values())
367
+ avail_numel = sum([p.numel() for p in zero_model_states[0].frozen_param_fragments.values()]) * world_size
368
+ print(f'Frozen params: Have {avail_numel} numels to process.')
369
+ print(f'Frozen params: Need {wanted_numel} numels in {wanted_params} params')
370
+
371
+ total_params = 0
372
+ total_numel = 0
373
+ for name, shape in zero_model_states[0].frozen_param_shapes.items():
374
+ total_params += 1
375
+ unpartitioned_numel = shape.numel()
376
+ total_numel += unpartitioned_numel
377
+
378
+ param_frags = tuple(model_state.frozen_param_fragments[name] for model_state in zero_model_states)
379
+ state_dict[name] = torch.cat(param_frags, 0).narrow(0, 0, unpartitioned_numel).view(shape)
380
+
381
+ partitioned_numel, partitioned_padding_numel = zero3_partitioned_param_info(unpartitioned_numel, world_size)
382
+
383
+ if debug:
384
+ print(
385
+ f"Frozen params: {total_params} {name} full shape: {shape} partition0 numel={partitioned_numel} partitioned_padding_numel={partitioned_padding_numel}"
386
+ )
387
+
388
+ print(f"Reconstructed Frozen fp32 state dict with {total_params} params {total_numel} elements")
389
+
390
+
391
+ class GatheredTensor:
392
+ """
393
+ A pseudo tensor that collects partitioned weights.
394
+ It is more memory efficient when there are multiple groups.
395
+ """
396
+
397
+ def __init__(self, flat_groups, flat_groups_offset, offset, partitioned_numel, shape):
398
+ self.flat_groups = flat_groups
399
+ self.flat_groups_offset = flat_groups_offset
400
+ self.offset = offset
401
+ self.partitioned_numel = partitioned_numel
402
+ self.shape = shape
403
+ self.dtype = self.flat_groups[0][0].dtype
404
+
405
+ def contiguous(self):
406
+ """
407
+ Merge partitioned weights from flat_groups into a single tensor.
408
+ """
409
+ end_idx = self.offset + self.partitioned_numel
410
+ world_size = len(self.flat_groups)
411
+ pad_flat_param_chunks = []
412
+
413
+ for rank_i in range(world_size):
414
+ # for each rank, we need to collect weights from related group/groups
415
+ flat_groups_at_rank_i = self.flat_groups[rank_i]
416
+ start_group_id = None
417
+ end_group_id = None
418
+ for group_id in range(len(self.flat_groups_offset)):
419
+ if self.flat_groups_offset[group_id] <= self.offset < self.flat_groups_offset[group_id + 1]:
420
+ start_group_id = group_id
421
+ if self.flat_groups_offset[group_id] < end_idx <= self.flat_groups_offset[group_id + 1]:
422
+ end_group_id = group_id
423
+ break
424
+ # collect weights from related group/groups
425
+ for group_id in range(start_group_id, end_group_id + 1):
426
+ flat_tensor = flat_groups_at_rank_i[group_id]
427
+ start_offset = self.offset - self.flat_groups_offset[group_id]
428
+ end_offset = min(end_idx, self.flat_groups_offset[group_id + 1]) - self.flat_groups_offset[group_id]
429
+ pad_flat_param_chunks.append(flat_tensor[start_offset:end_offset])
430
+
431
+ # collect weights from all ranks
432
+ pad_flat_param = torch.cat(pad_flat_param_chunks, dim=0)
433
+ param = pad_flat_param[:self.shape.numel()].view(self.shape).contiguous()
434
+ return param
435
+
436
+
437
+ def _zero3_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states):
438
+ param_shapes = zero_model_states[0].param_shapes
439
+ avail_numel = sum([flat_group.numel() for flat_group in fp32_flat_groups[0]]) * world_size
440
+
441
+ # Reconstruction protocol: For zero3 we need to zip the partitions together at boundary of each
442
+ # param, re-consolidating each param, while dealing with padding if any
443
+
444
+ # merge list of dicts, preserving order
445
+ param_shapes = {k: v for d in param_shapes for k, v in d.items()}
446
+
447
+ if debug:
448
+ for i in range(world_size):
449
+ print(f"{FP32_FLAT_GROUPS}[{i}].shape={fp32_flat_groups[i].shape}")
450
+
451
+ wanted_params = len(param_shapes)
452
+ wanted_numel = sum(shape.numel() for shape in param_shapes.values())
453
+ # not asserting if there is a mismatch due to possible padding
454
+ avail_numel = fp32_flat_groups[0].numel() * world_size
455
+ print(f"Trainable params: Have {avail_numel} numels to process.")
456
+ print(f"Trainable params: Need {wanted_numel} numels in {wanted_params} params.")
457
+
458
+ # params
459
+ # XXX: for huge models that can't fit into the host's RAM we will have to recode this to support
460
+ # out-of-core computing solution
461
+ offset = 0
462
+ total_numel = 0
463
+ total_params = 0
464
+ flat_groups_offset = [0] + list(np.cumsum([flat_tensor.numel() for flat_tensor in fp32_flat_groups[0]]))
465
+ for name, shape in tqdm(param_shapes.items(), desc='Gathering sharded weights'):
466
+ unpartitioned_numel = shape.numel()
467
+ total_numel += unpartitioned_numel
468
+ total_params += 1
469
+ partitioned_numel, partitioned_padding_numel = zero3_partitioned_param_info(unpartitioned_numel, world_size)
470
+
471
+ if debug:
472
+ print(
473
+ f"Trainable params: {total_params} {name} full shape: {shape} partition0 numel={partitioned_numel} partitioned_padding_numel={partitioned_padding_numel}"
474
+ )
475
+
476
+ # memory efficient tensor
477
+ tensor = GatheredTensor(fp32_flat_groups, flat_groups_offset, offset, partitioned_numel, shape)
478
+ state_dict[name] = tensor
479
+ offset += partitioned_numel
480
+
481
+ offset *= world_size
482
+
483
+ # Sanity check
484
+ if offset != avail_numel:
485
+ raise ValueError(f"consumed {offset} numels out of {avail_numel} - something is wrong")
486
+
487
+ print(f"Reconstructed Trainable fp32 state dict with {total_params} params {total_numel} elements")
488
+
489
+
490
+ def _get_fp32_state_dict_from_zero3_checkpoint(world_size, fp32_flat_groups, zero_model_states,
491
+ exclude_frozen_parameters):
492
+ state_dict = OrderedDict()
493
+
494
+ # buffers
495
+ buffers = zero_model_states[0].buffers
496
+ state_dict.update(buffers)
497
+ if debug:
498
+ print(f"added {len(buffers)} buffers")
499
+
500
+ if not exclude_frozen_parameters:
501
+ _zero3_merge_frozen_params(state_dict, world_size, zero_model_states)
502
+
503
+ _zero3_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states)
504
+
505
+ # recover shared parameters
506
+ for pair in zero_model_states[0].shared_params:
507
+ if pair[1] in state_dict:
508
+ state_dict[pair[0]] = state_dict[pair[1]]
509
+
510
+ return state_dict
511
+
512
+
513
+ def to_torch_tensor(state_dict, return_empty_tensor=False):
514
+ """
515
+ Convert state_dict of GatheredTensor to torch tensor
516
+ """
517
+ torch_state_dict = {}
518
+ converted_tensors = {}
519
+ for name, tensor in state_dict.items():
520
+ tensor_id = id(tensor)
521
+ if tensor_id in converted_tensors: # shared tensors
522
+ shared_tensor = torch_state_dict[converted_tensors[tensor_id]]
523
+ torch_state_dict[name] = shared_tensor
524
+ else:
525
+ converted_tensors[tensor_id] = name
526
+ if return_empty_tensor:
527
+ torch_state_dict[name] = torch.empty(tensor.shape, dtype=tensor.dtype)
528
+ else:
529
+ torch_state_dict[name] = tensor.contiguous()
530
+ return torch_state_dict
531
+
532
+
533
+ def get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir,
534
+ tag=None,
535
+ exclude_frozen_parameters=False,
536
+ lazy_mode=False):
537
+ """
538
+ Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated state_dict that can be loaded with
539
+ ``load_state_dict()`` and used for training without DeepSpeed or shared with others, for example
540
+ via a model hub.
541
+
542
+ Args:
543
+ - ``checkpoint_dir``: path to the desired checkpoint folder
544
+ - ``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``
545
+ - ``exclude_frozen_parameters``: exclude frozen parameters
546
+ - ``lazy_mode``: get state_dict in lazy mode. It returns a dict of pesduo tensor instead of torch tensor, which is more memory efficient.
547
+ Convert the pesduo tensor to torch tensor by ``.contiguous()``
548
+
549
+ Returns:
550
+ - pytorch ``state_dict``
551
+
552
+ A typical usage might be ::
553
+
554
+ from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint
555
+ # do the training and checkpoint saving
556
+ state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir) # already on cpu
557
+ model = model.cpu() # move to cpu
558
+ model.load_state_dict(state_dict)
559
+ # submit to model hub or save the model to share with others
560
+
561
+ In this example the ``model`` will no longer be usable in the deepspeed context of the same
562
+ application. i.e. you will need to re-initialize the deepspeed engine, since
563
+ ``model.load_state_dict(state_dict)`` will remove all the deepspeed magic from it.
564
+
565
+ If you want it all done for you, use ``load_state_dict_from_zero_checkpoint`` instead.
566
+
567
+ Note: the above usage may not work if your application doesn't have sufficient free CPU memory.
568
+ You may need to use the offline approach using the ``zero_to_fp32.py`` script that is saved with
569
+ the checkpoint. Or you can load state_dict in lazy mode ::
570
+
571
+ from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint
572
+ state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, lazy_mode=True) # not on cpu
573
+ for name, lazy_tensor in state_dict.item():
574
+ tensor = lazy_tensor.contiguous() # to cpu
575
+ print(name, tensor)
576
+ # del tensor to release memory if it no longer in use
577
+ """
578
+ if tag is None:
579
+ latest_path = os.path.join(checkpoint_dir, 'latest')
580
+ if os.path.isfile(latest_path):
581
+ with open(latest_path, 'r') as fd:
582
+ tag = fd.read().strip()
583
+ else:
584
+ raise ValueError(f"Unable to find 'latest' file at {latest_path}")
585
+
586
+ ds_checkpoint_dir = os.path.join(checkpoint_dir, tag)
587
+
588
+ if not os.path.isdir(ds_checkpoint_dir):
589
+ raise FileNotFoundError(f"Directory '{ds_checkpoint_dir}' doesn't exist")
590
+
591
+ state_dict = _get_fp32_state_dict_from_zero_checkpoint(ds_checkpoint_dir, exclude_frozen_parameters)
592
+ if lazy_mode:
593
+ return state_dict
594
+ else:
595
+ return to_torch_tensor(state_dict)
596
+
597
+
598
+ def convert_zero_checkpoint_to_fp32_state_dict(checkpoint_dir,
599
+ output_dir,
600
+ max_shard_size="5GB",
601
+ safe_serialization=False,
602
+ tag=None,
603
+ exclude_frozen_parameters=False):
604
+ """
605
+ Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated ``state_dict`` file that can be
606
+ loaded with ``torch.load(file)`` + ``load_state_dict()`` and used for training without DeepSpeed.
607
+
608
+ Args:
609
+ - ``checkpoint_dir``: path to the desired checkpoint folder. (one that contains the tag-folder, like ``global_step14``)
610
+ - ``output_dir``: directory to the pytorch fp32 state_dict output files
611
+ - ``max_shard_size``: the maximum size for a checkpoint before being sharded, default value is 5GB
612
+ - ``safe_serialization``: whether to save the model using `safetensors` or the traditional PyTorch way (that uses `pickle`).
613
+ - ``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``
614
+ - ``exclude_frozen_parameters``: exclude frozen parameters
615
+ """
616
+
617
+ # Dependency pre-check
618
+ if safe_serialization:
619
+ try:
620
+ from safetensors.torch import save_file
621
+ except ImportError:
622
+ print('If you want to use `safe_serialization`, please `pip install safetensors`')
623
+ raise
624
+ if max_shard_size is not None:
625
+ try:
626
+ from huggingface_hub import split_torch_state_dict_into_shards
627
+ except ImportError:
628
+ print('If you want to use `max_shard_size`, please `pip install huggingface_hub`')
629
+ raise
630
+
631
+ # Convert zero checkpoint to state_dict
632
+ state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir,
633
+ tag,
634
+ exclude_frozen_parameters,
635
+ lazy_mode=True)
636
+
637
+ # Shard the model if it is too big.
638
+ weights_name = "model.safetensors" if safe_serialization else "pytorch_model.bin"
639
+ if max_shard_size is not None:
640
+ filename_pattern = weights_name.replace(".bin", "{suffix}.bin").replace(".safetensors", "{suffix}.safetensors")
641
+ # an memory-efficient approach for sharding
642
+ empty_state_dict = to_torch_tensor(state_dict, return_empty_tensor=True)
643
+ state_dict_split = split_torch_state_dict_into_shards(empty_state_dict,
644
+ filename_pattern=filename_pattern,
645
+ max_shard_size=max_shard_size)
646
+ else:
647
+ from collections import namedtuple
648
+ StateDictSplit = namedtuple("StateDictSplit", ["is_sharded", "filename_to_tensors"])
649
+ state_dict_split = StateDictSplit(is_sharded=False,
650
+ filename_to_tensors={weights_name: list(state_dict.keys())})
651
+
652
+ # Save the model by shard
653
+ os.makedirs(output_dir, exist_ok=True)
654
+ filename_to_tensors = state_dict_split.filename_to_tensors.items()
655
+ for shard_file, tensors in tqdm(filename_to_tensors, desc="Saving checkpoint shards"):
656
+ shard_state_dict = {tensor_name: state_dict[tensor_name] for tensor_name in tensors}
657
+ shard_state_dict = to_torch_tensor(shard_state_dict)
658
+ output_path = os.path.join(output_dir, shard_file)
659
+ if safe_serialization:
660
+ save_file(shard_state_dict, output_path, metadata={"format": "pt"})
661
+ else:
662
+ torch.save(shard_state_dict, output_path)
663
+ # release the memory of current shard
664
+ for tensor_name in list(shard_state_dict.keys()):
665
+ del state_dict[tensor_name]
666
+ del shard_state_dict[tensor_name]
667
+ del shard_state_dict
668
+ gc.collect()
669
+
670
+ # Save index if sharded
671
+ if state_dict_split.is_sharded:
672
+ index = {
673
+ "metadata": state_dict_split.metadata,
674
+ "weight_map": state_dict_split.tensor_to_filename,
675
+ }
676
+ save_index_file = "model.safetensors.index.json" if safe_serialization else "pytorch_model.bin.index.json"
677
+ save_index_file = os.path.join(output_dir, save_index_file)
678
+ with open(save_index_file, "w", encoding="utf-8") as f:
679
+ content = json.dumps(index, indent=2, sort_keys=True) + "\n"
680
+ f.write(content)
681
+
682
+
683
+ def load_state_dict_from_zero_checkpoint(model, checkpoint_dir, tag=None):
684
+ """
685
+ 1. Put the provided model to cpu
686
+ 2. Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated ``state_dict``
687
+ 3. Load it into the provided model
688
+
689
+ Args:
690
+ - ``model``: the model object to update
691
+ - ``checkpoint_dir``: path to the desired checkpoint folder. (one that contains the tag-folder, like ``global_step14``)
692
+ - ``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``
693
+
694
+ Returns:
695
+ - ``model`: modified model
696
+
697
+ Make sure you have plenty of CPU memory available before you call this function. If you don't
698
+ have enough use the ``zero_to_fp32.py`` utility to do the conversion. You will find it
699
+ conveniently placed for you in the checkpoint folder.
700
+
701
+ A typical usage might be ::
702
+
703
+ from deepspeed.utils.zero_to_fp32 import load_state_dict_from_zero_checkpoint
704
+ model = load_state_dict_from_zero_checkpoint(trainer.model, checkpoint_dir)
705
+ # submit to model hub or save the model to share with others
706
+
707
+ Note, that once this was run, the ``model`` will no longer be usable in the deepspeed context
708
+ of the same application. i.e. you will need to re-initialize the deepspeed engine, since
709
+ ``model.load_state_dict(state_dict)`` will remove all the deepspeed magic from it.
710
+
711
+ """
712
+ logger.info("Extracting fp32 weights")
713
+ state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, tag)
714
+
715
+ logger.info("Overwriting model with fp32 weights")
716
+ model = model.cpu()
717
+ model.load_state_dict(state_dict, strict=False)
718
+
719
+ return model
720
+
721
+
722
+ if __name__ == "__main__":
723
+ parser = argparse.ArgumentParser()
724
+ parser.add_argument("checkpoint_dir",
725
+ type=str,
726
+ help="path to the desired checkpoint folder, e.g., path/checkpoint-12")
727
+ parser.add_argument("output_dir",
728
+ type=str,
729
+ help="directory to the pytorch fp32 state_dict output files"
730
+ "(e.g. path/checkpoint-12-output/)")
731
+ parser.add_argument(
732
+ "--max_shard_size",
733
+ type=str,
734
+ default="5GB",
735
+ help="The maximum size for a checkpoint before being sharded. Checkpoints shard will then be each of size"
736
+ "lower than this size. If expressed as a string, needs to be digits followed by a unit (like `5MB`"
737
+ "We default it to 5GB in order for models to be able to run easily on free-tier google colab instances"
738
+ "without CPU OOM issues.")
739
+ parser.add_argument(
740
+ "--safe_serialization",
741
+ default=False,
742
+ action='store_true',
743
+ help="Whether to save the model using `safetensors` or the traditional PyTorch way (that uses `pickle`).")
744
+ parser.add_argument("-t",
745
+ "--tag",
746
+ type=str,
747
+ default=None,
748
+ help="checkpoint tag used as a unique identifier for checkpoint. e.g., global_step1")
749
+ parser.add_argument("--exclude_frozen_parameters", action='store_true', help="exclude frozen parameters")
750
+ parser.add_argument("-d", "--debug", action='store_true', help="enable debug")
751
+ args = parser.parse_args()
752
+
753
+ debug = args.debug
754
+
755
+ convert_zero_checkpoint_to_fp32_state_dict(args.checkpoint_dir,
756
+ args.output_dir,
757
+ max_shard_size=args.max_shard_size,
758
+ safe_serialization=args.safe_serialization,
759
+ tag=args.tag,
760
+ exclude_frozen_parameters=args.exclude_frozen_parameters)
openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1230/config.json ADDED
@@ -0,0 +1,39 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "architectures": [
3
+ "LlamaForCausalLM"
4
+ ],
5
+ "attention_bias": false,
6
+ "attention_dropout": 0.0,
7
+ "bos_token_id": 128000,
8
+ "dtype": "bfloat16",
9
+ "eos_token_id": [
10
+ 128001,
11
+ 128008,
12
+ 128009
13
+ ],
14
+ "head_dim": 64,
15
+ "hidden_act": "silu",
16
+ "hidden_size": 2048,
17
+ "initializer_range": 0.02,
18
+ "intermediate_size": 8192,
19
+ "max_position_embeddings": 131072,
20
+ "mlp_bias": false,
21
+ "model_type": "llama",
22
+ "num_attention_heads": 32,
23
+ "num_hidden_layers": 16,
24
+ "num_key_value_heads": 8,
25
+ "pretraining_tp": 1,
26
+ "rms_norm_eps": 1e-05,
27
+ "rope_scaling": {
28
+ "factor": 32.0,
29
+ "high_freq_factor": 4.0,
30
+ "low_freq_factor": 1.0,
31
+ "original_max_position_embeddings": 8192,
32
+ "rope_type": "llama3"
33
+ },
34
+ "rope_theta": 500000.0,
35
+ "tie_word_embeddings": true,
36
+ "transformers_version": "4.57.1",
37
+ "use_cache": false,
38
+ "vocab_size": 128256
39
+ }
openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1230/generation_config.json ADDED
@@ -0,0 +1,12 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "bos_token_id": 128000,
3
+ "do_sample": true,
4
+ "eos_token_id": [
5
+ 128001,
6
+ 128008,
7
+ 128009
8
+ ],
9
+ "temperature": 0.6,
10
+ "top_p": 0.9,
11
+ "transformers_version": "4.57.1"
12
+ }
openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1230/latest ADDED
@@ -0,0 +1 @@
 
 
1
+ global_step1230
openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1230/model.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:a48d40b419c750e901d492309f615ed524a2647afd9b8afb05d5516ce0df53b6
3
+ size 2996982344
openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1230/scheduler.pt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:e7e8477210f775a7c1c3ea7cbcf76343d146771b5b05309a94b7d69ce3332955
3
+ size 1465
openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1230/trainer_state.json ADDED
@@ -0,0 +1,895 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "best_global_step": null,
3
+ "best_metric": null,
4
+ "best_model_checkpoint": null,
5
+ "epoch": 1.4984765386959171,
6
+ "eval_steps": 500,
7
+ "global_step": 1230,
8
+ "is_hyper_param_search": false,
9
+ "is_local_process_zero": true,
10
+ "is_world_process_zero": true,
11
+ "log_history": [
12
+ {
13
+ "epoch": 0.01218769043266301,
14
+ "grad_norm": 0.36486390233039856,
15
+ "learning_rate": 5e-05,
16
+ "loss": 1.2213,
17
+ "step": 10
18
+ },
19
+ {
20
+ "epoch": 0.02437538086532602,
21
+ "grad_norm": 0.2953564524650574,
22
+ "learning_rate": 5e-05,
23
+ "loss": 1.0223,
24
+ "step": 20
25
+ },
26
+ {
27
+ "epoch": 0.03656307129798903,
28
+ "grad_norm": 0.2766735851764679,
29
+ "learning_rate": 5e-05,
30
+ "loss": 0.9589,
31
+ "step": 30
32
+ },
33
+ {
34
+ "epoch": 0.04875076173065204,
35
+ "grad_norm": 0.32770997285842896,
36
+ "learning_rate": 5e-05,
37
+ "loss": 0.9081,
38
+ "step": 40
39
+ },
40
+ {
41
+ "epoch": 0.06093845216331505,
42
+ "grad_norm": 0.3052276372909546,
43
+ "learning_rate": 5e-05,
44
+ "loss": 0.8752,
45
+ "step": 50
46
+ },
47
+ {
48
+ "epoch": 0.07312614259597806,
49
+ "grad_norm": 0.240146666765213,
50
+ "learning_rate": 5e-05,
51
+ "loss": 0.873,
52
+ "step": 60
53
+ },
54
+ {
55
+ "epoch": 0.08531383302864107,
56
+ "grad_norm": 0.23872187733650208,
57
+ "learning_rate": 5e-05,
58
+ "loss": 0.8587,
59
+ "step": 70
60
+ },
61
+ {
62
+ "epoch": 0.09750152346130408,
63
+ "grad_norm": 0.32638487219810486,
64
+ "learning_rate": 5e-05,
65
+ "loss": 0.8656,
66
+ "step": 80
67
+ },
68
+ {
69
+ "epoch": 0.10968921389396709,
70
+ "grad_norm": 0.25019028782844543,
71
+ "learning_rate": 5e-05,
72
+ "loss": 0.8685,
73
+ "step": 90
74
+ },
75
+ {
76
+ "epoch": 0.1218769043266301,
77
+ "grad_norm": 0.2507675886154175,
78
+ "learning_rate": 5e-05,
79
+ "loss": 0.8662,
80
+ "step": 100
81
+ },
82
+ {
83
+ "epoch": 0.1340645947592931,
84
+ "grad_norm": 0.27966755628585815,
85
+ "learning_rate": 5e-05,
86
+ "loss": 0.8428,
87
+ "step": 110
88
+ },
89
+ {
90
+ "epoch": 0.14625228519195613,
91
+ "grad_norm": 0.2622673213481903,
92
+ "learning_rate": 5e-05,
93
+ "loss": 0.8416,
94
+ "step": 120
95
+ },
96
+ {
97
+ "epoch": 0.15843997562461914,
98
+ "grad_norm": 0.3003966510295868,
99
+ "learning_rate": 5e-05,
100
+ "loss": 0.8528,
101
+ "step": 130
102
+ },
103
+ {
104
+ "epoch": 0.17062766605728213,
105
+ "grad_norm": 0.32308319211006165,
106
+ "learning_rate": 5e-05,
107
+ "loss": 0.8357,
108
+ "step": 140
109
+ },
110
+ {
111
+ "epoch": 0.18281535648994515,
112
+ "grad_norm": 0.27279096841812134,
113
+ "learning_rate": 5e-05,
114
+ "loss": 0.8432,
115
+ "step": 150
116
+ },
117
+ {
118
+ "epoch": 0.19500304692260817,
119
+ "grad_norm": 0.27976104617118835,
120
+ "learning_rate": 5e-05,
121
+ "loss": 0.8119,
122
+ "step": 160
123
+ },
124
+ {
125
+ "epoch": 0.2071907373552712,
126
+ "grad_norm": 0.290866881608963,
127
+ "learning_rate": 5e-05,
128
+ "loss": 0.8096,
129
+ "step": 170
130
+ },
131
+ {
132
+ "epoch": 0.21937842778793418,
133
+ "grad_norm": 0.29430949687957764,
134
+ "learning_rate": 5e-05,
135
+ "loss": 0.819,
136
+ "step": 180
137
+ },
138
+ {
139
+ "epoch": 0.2315661182205972,
140
+ "grad_norm": 0.27714136242866516,
141
+ "learning_rate": 5e-05,
142
+ "loss": 0.8187,
143
+ "step": 190
144
+ },
145
+ {
146
+ "epoch": 0.2437538086532602,
147
+ "grad_norm": 0.2997029423713684,
148
+ "learning_rate": 5e-05,
149
+ "loss": 0.808,
150
+ "step": 200
151
+ },
152
+ {
153
+ "epoch": 0.25594149908592323,
154
+ "grad_norm": 0.2539573609828949,
155
+ "learning_rate": 5e-05,
156
+ "loss": 0.7943,
157
+ "step": 210
158
+ },
159
+ {
160
+ "epoch": 0.2681291895185862,
161
+ "grad_norm": 0.23294171690940857,
162
+ "learning_rate": 5e-05,
163
+ "loss": 0.8188,
164
+ "step": 220
165
+ },
166
+ {
167
+ "epoch": 0.28031687995124926,
168
+ "grad_norm": 0.25016292929649353,
169
+ "learning_rate": 5e-05,
170
+ "loss": 0.8035,
171
+ "step": 230
172
+ },
173
+ {
174
+ "epoch": 0.29250457038391225,
175
+ "grad_norm": 0.37125322222709656,
176
+ "learning_rate": 5e-05,
177
+ "loss": 0.785,
178
+ "step": 240
179
+ },
180
+ {
181
+ "epoch": 0.30469226081657524,
182
+ "grad_norm": 0.21520452201366425,
183
+ "learning_rate": 5e-05,
184
+ "loss": 0.8088,
185
+ "step": 250
186
+ },
187
+ {
188
+ "epoch": 0.3168799512492383,
189
+ "grad_norm": 0.20990602672100067,
190
+ "learning_rate": 5e-05,
191
+ "loss": 0.7813,
192
+ "step": 260
193
+ },
194
+ {
195
+ "epoch": 0.3290676416819013,
196
+ "grad_norm": 0.24193961918354034,
197
+ "learning_rate": 5e-05,
198
+ "loss": 0.7979,
199
+ "step": 270
200
+ },
201
+ {
202
+ "epoch": 0.34125533211456427,
203
+ "grad_norm": 0.22766050696372986,
204
+ "learning_rate": 5e-05,
205
+ "loss": 0.785,
206
+ "step": 280
207
+ },
208
+ {
209
+ "epoch": 0.3534430225472273,
210
+ "grad_norm": 0.2662702202796936,
211
+ "learning_rate": 5e-05,
212
+ "loss": 0.7766,
213
+ "step": 290
214
+ },
215
+ {
216
+ "epoch": 0.3656307129798903,
217
+ "grad_norm": 0.23020495474338531,
218
+ "learning_rate": 5e-05,
219
+ "loss": 0.7687,
220
+ "step": 300
221
+ },
222
+ {
223
+ "epoch": 0.3778184034125533,
224
+ "grad_norm": 0.21770520508289337,
225
+ "learning_rate": 5e-05,
226
+ "loss": 0.7815,
227
+ "step": 310
228
+ },
229
+ {
230
+ "epoch": 0.39000609384521634,
231
+ "grad_norm": 0.22677403688430786,
232
+ "learning_rate": 5e-05,
233
+ "loss": 0.7999,
234
+ "step": 320
235
+ },
236
+ {
237
+ "epoch": 0.40219378427787933,
238
+ "grad_norm": 0.24891860783100128,
239
+ "learning_rate": 5e-05,
240
+ "loss": 0.7946,
241
+ "step": 330
242
+ },
243
+ {
244
+ "epoch": 0.4143814747105424,
245
+ "grad_norm": 0.2506854832172394,
246
+ "learning_rate": 5e-05,
247
+ "loss": 0.7889,
248
+ "step": 340
249
+ },
250
+ {
251
+ "epoch": 0.42656916514320536,
252
+ "grad_norm": 0.24388180673122406,
253
+ "learning_rate": 5e-05,
254
+ "loss": 0.79,
255
+ "step": 350
256
+ },
257
+ {
258
+ "epoch": 0.43875685557586835,
259
+ "grad_norm": 0.22329650819301605,
260
+ "learning_rate": 5e-05,
261
+ "loss": 0.7632,
262
+ "step": 360
263
+ },
264
+ {
265
+ "epoch": 0.4509445460085314,
266
+ "grad_norm": 0.2240927815437317,
267
+ "learning_rate": 5e-05,
268
+ "loss": 0.8042,
269
+ "step": 370
270
+ },
271
+ {
272
+ "epoch": 0.4631322364411944,
273
+ "grad_norm": 0.2578203082084656,
274
+ "learning_rate": 5e-05,
275
+ "loss": 0.7667,
276
+ "step": 380
277
+ },
278
+ {
279
+ "epoch": 0.4753199268738574,
280
+ "grad_norm": 0.25349217653274536,
281
+ "learning_rate": 5e-05,
282
+ "loss": 0.7911,
283
+ "step": 390
284
+ },
285
+ {
286
+ "epoch": 0.4875076173065204,
287
+ "grad_norm": 0.254446804523468,
288
+ "learning_rate": 5e-05,
289
+ "loss": 0.7909,
290
+ "step": 400
291
+ },
292
+ {
293
+ "epoch": 0.4996953077391834,
294
+ "grad_norm": 0.24253171682357788,
295
+ "learning_rate": 5e-05,
296
+ "loss": 0.7896,
297
+ "step": 410
298
+ },
299
+ {
300
+ "epoch": 0.5118829981718465,
301
+ "grad_norm": 0.27901458740234375,
302
+ "learning_rate": 5e-05,
303
+ "loss": 0.7828,
304
+ "step": 420
305
+ },
306
+ {
307
+ "epoch": 0.5240706886045094,
308
+ "grad_norm": 0.31884434819221497,
309
+ "learning_rate": 5e-05,
310
+ "loss": 0.7919,
311
+ "step": 430
312
+ },
313
+ {
314
+ "epoch": 0.5362583790371724,
315
+ "grad_norm": 0.22920799255371094,
316
+ "learning_rate": 5e-05,
317
+ "loss": 0.7777,
318
+ "step": 440
319
+ },
320
+ {
321
+ "epoch": 0.5484460694698354,
322
+ "grad_norm": 0.22606618702411652,
323
+ "learning_rate": 5e-05,
324
+ "loss": 0.7661,
325
+ "step": 450
326
+ },
327
+ {
328
+ "epoch": 0.5606337599024985,
329
+ "grad_norm": 0.21604052186012268,
330
+ "learning_rate": 5e-05,
331
+ "loss": 0.769,
332
+ "step": 460
333
+ },
334
+ {
335
+ "epoch": 0.5728214503351615,
336
+ "grad_norm": 0.2296101152896881,
337
+ "learning_rate": 5e-05,
338
+ "loss": 0.7626,
339
+ "step": 470
340
+ },
341
+ {
342
+ "epoch": 0.5850091407678245,
343
+ "grad_norm": 0.24162150919437408,
344
+ "learning_rate": 5e-05,
345
+ "loss": 0.7831,
346
+ "step": 480
347
+ },
348
+ {
349
+ "epoch": 0.5971968312004875,
350
+ "grad_norm": 0.22364799678325653,
351
+ "learning_rate": 5e-05,
352
+ "loss": 0.7737,
353
+ "step": 490
354
+ },
355
+ {
356
+ "epoch": 0.6093845216331505,
357
+ "grad_norm": 0.23953403532505035,
358
+ "learning_rate": 5e-05,
359
+ "loss": 0.787,
360
+ "step": 500
361
+ },
362
+ {
363
+ "epoch": 0.6215722120658135,
364
+ "grad_norm": 0.22349311411380768,
365
+ "learning_rate": 5e-05,
366
+ "loss": 0.7763,
367
+ "step": 510
368
+ },
369
+ {
370
+ "epoch": 0.6337599024984766,
371
+ "grad_norm": 0.21996377408504486,
372
+ "learning_rate": 5e-05,
373
+ "loss": 0.7596,
374
+ "step": 520
375
+ },
376
+ {
377
+ "epoch": 0.6459475929311396,
378
+ "grad_norm": 0.245538130402565,
379
+ "learning_rate": 5e-05,
380
+ "loss": 0.7785,
381
+ "step": 530
382
+ },
383
+ {
384
+ "epoch": 0.6581352833638026,
385
+ "grad_norm": 0.23442520201206207,
386
+ "learning_rate": 5e-05,
387
+ "loss": 0.7633,
388
+ "step": 540
389
+ },
390
+ {
391
+ "epoch": 0.6703229737964655,
392
+ "grad_norm": 0.22274106740951538,
393
+ "learning_rate": 5e-05,
394
+ "loss": 0.7756,
395
+ "step": 550
396
+ },
397
+ {
398
+ "epoch": 0.6825106642291285,
399
+ "grad_norm": 0.21717530488967896,
400
+ "learning_rate": 5e-05,
401
+ "loss": 0.7726,
402
+ "step": 560
403
+ },
404
+ {
405
+ "epoch": 0.6946983546617916,
406
+ "grad_norm": 0.2465955913066864,
407
+ "learning_rate": 5e-05,
408
+ "loss": 0.7476,
409
+ "step": 570
410
+ },
411
+ {
412
+ "epoch": 0.7068860450944546,
413
+ "grad_norm": 0.2318575084209442,
414
+ "learning_rate": 5e-05,
415
+ "loss": 0.7584,
416
+ "step": 580
417
+ },
418
+ {
419
+ "epoch": 0.7190737355271176,
420
+ "grad_norm": 0.25114378333091736,
421
+ "learning_rate": 5e-05,
422
+ "loss": 0.7625,
423
+ "step": 590
424
+ },
425
+ {
426
+ "epoch": 0.7312614259597806,
427
+ "grad_norm": 0.19145755469799042,
428
+ "learning_rate": 5e-05,
429
+ "loss": 0.7577,
430
+ "step": 600
431
+ },
432
+ {
433
+ "epoch": 0.7434491163924436,
434
+ "grad_norm": 0.19129066169261932,
435
+ "learning_rate": 5e-05,
436
+ "loss": 0.7772,
437
+ "step": 610
438
+ },
439
+ {
440
+ "epoch": 0.7556368068251066,
441
+ "grad_norm": 0.2014475166797638,
442
+ "learning_rate": 5e-05,
443
+ "loss": 0.7601,
444
+ "step": 620
445
+ },
446
+ {
447
+ "epoch": 0.7678244972577697,
448
+ "grad_norm": 0.21480509638786316,
449
+ "learning_rate": 5e-05,
450
+ "loss": 0.7665,
451
+ "step": 630
452
+ },
453
+ {
454
+ "epoch": 0.7800121876904327,
455
+ "grad_norm": 0.2633228302001953,
456
+ "learning_rate": 5e-05,
457
+ "loss": 0.7565,
458
+ "step": 640
459
+ },
460
+ {
461
+ "epoch": 0.7921998781230957,
462
+ "grad_norm": 0.24633203446865082,
463
+ "learning_rate": 5e-05,
464
+ "loss": 0.7602,
465
+ "step": 650
466
+ },
467
+ {
468
+ "epoch": 0.8043875685557587,
469
+ "grad_norm": 0.19770534336566925,
470
+ "learning_rate": 5e-05,
471
+ "loss": 0.7685,
472
+ "step": 660
473
+ },
474
+ {
475
+ "epoch": 0.8165752589884216,
476
+ "grad_norm": 0.2267930805683136,
477
+ "learning_rate": 5e-05,
478
+ "loss": 0.7516,
479
+ "step": 670
480
+ },
481
+ {
482
+ "epoch": 0.8287629494210847,
483
+ "grad_norm": 0.238111674785614,
484
+ "learning_rate": 5e-05,
485
+ "loss": 0.7487,
486
+ "step": 680
487
+ },
488
+ {
489
+ "epoch": 0.8409506398537477,
490
+ "grad_norm": 0.19528904557228088,
491
+ "learning_rate": 5e-05,
492
+ "loss": 0.7467,
493
+ "step": 690
494
+ },
495
+ {
496
+ "epoch": 0.8531383302864107,
497
+ "grad_norm": 0.21305914223194122,
498
+ "learning_rate": 5e-05,
499
+ "loss": 0.7527,
500
+ "step": 700
501
+ },
502
+ {
503
+ "epoch": 0.8653260207190737,
504
+ "grad_norm": 0.2205205112695694,
505
+ "learning_rate": 5e-05,
506
+ "loss": 0.7639,
507
+ "step": 710
508
+ },
509
+ {
510
+ "epoch": 0.8775137111517367,
511
+ "grad_norm": 0.22429127991199493,
512
+ "learning_rate": 5e-05,
513
+ "loss": 0.7744,
514
+ "step": 720
515
+ },
516
+ {
517
+ "epoch": 0.8897014015843998,
518
+ "grad_norm": 0.23755772411823273,
519
+ "learning_rate": 5e-05,
520
+ "loss": 0.7446,
521
+ "step": 730
522
+ },
523
+ {
524
+ "epoch": 0.9018890920170628,
525
+ "grad_norm": 0.17592693865299225,
526
+ "learning_rate": 5e-05,
527
+ "loss": 0.7442,
528
+ "step": 740
529
+ },
530
+ {
531
+ "epoch": 0.9140767824497258,
532
+ "grad_norm": 0.2266898900270462,
533
+ "learning_rate": 5e-05,
534
+ "loss": 0.7326,
535
+ "step": 750
536
+ },
537
+ {
538
+ "epoch": 0.9262644728823888,
539
+ "grad_norm": 0.23541569709777832,
540
+ "learning_rate": 5e-05,
541
+ "loss": 0.7338,
542
+ "step": 760
543
+ },
544
+ {
545
+ "epoch": 0.9384521633150518,
546
+ "grad_norm": 0.22233888506889343,
547
+ "learning_rate": 5e-05,
548
+ "loss": 0.7385,
549
+ "step": 770
550
+ },
551
+ {
552
+ "epoch": 0.9506398537477148,
553
+ "grad_norm": 0.23714767396450043,
554
+ "learning_rate": 5e-05,
555
+ "loss": 0.7306,
556
+ "step": 780
557
+ },
558
+ {
559
+ "epoch": 0.9628275441803779,
560
+ "grad_norm": 0.21191060543060303,
561
+ "learning_rate": 5e-05,
562
+ "loss": 0.736,
563
+ "step": 790
564
+ },
565
+ {
566
+ "epoch": 0.9750152346130408,
567
+ "grad_norm": 0.22759106755256653,
568
+ "learning_rate": 5e-05,
569
+ "loss": 0.7505,
570
+ "step": 800
571
+ },
572
+ {
573
+ "epoch": 0.9872029250457038,
574
+ "grad_norm": 0.19895903766155243,
575
+ "learning_rate": 5e-05,
576
+ "loss": 0.7536,
577
+ "step": 810
578
+ },
579
+ {
580
+ "epoch": 0.9993906154783668,
581
+ "grad_norm": 0.3061685860157013,
582
+ "learning_rate": 5e-05,
583
+ "loss": 0.7493,
584
+ "step": 820
585
+ },
586
+ {
587
+ "epoch": 1.0109689213893966,
588
+ "grad_norm": 0.18868015706539154,
589
+ "learning_rate": 5e-05,
590
+ "loss": 0.7342,
591
+ "step": 830
592
+ },
593
+ {
594
+ "epoch": 1.0231566118220596,
595
+ "grad_norm": 0.1797628402709961,
596
+ "learning_rate": 5e-05,
597
+ "loss": 0.7192,
598
+ "step": 840
599
+ },
600
+ {
601
+ "epoch": 1.0353443022547226,
602
+ "grad_norm": 0.19511596858501434,
603
+ "learning_rate": 5e-05,
604
+ "loss": 0.7002,
605
+ "step": 850
606
+ },
607
+ {
608
+ "epoch": 1.0475319926873858,
609
+ "grad_norm": 0.2123185396194458,
610
+ "learning_rate": 5e-05,
611
+ "loss": 0.7236,
612
+ "step": 860
613
+ },
614
+ {
615
+ "epoch": 1.0597196831200488,
616
+ "grad_norm": 0.2047470211982727,
617
+ "learning_rate": 5e-05,
618
+ "loss": 0.711,
619
+ "step": 870
620
+ },
621
+ {
622
+ "epoch": 1.0719073735527118,
623
+ "grad_norm": 0.18809868395328522,
624
+ "learning_rate": 5e-05,
625
+ "loss": 0.7253,
626
+ "step": 880
627
+ },
628
+ {
629
+ "epoch": 1.0840950639853748,
630
+ "grad_norm": 0.2207012176513672,
631
+ "learning_rate": 5e-05,
632
+ "loss": 0.6971,
633
+ "step": 890
634
+ },
635
+ {
636
+ "epoch": 1.0962827544180378,
637
+ "grad_norm": 0.1912035495042801,
638
+ "learning_rate": 5e-05,
639
+ "loss": 0.7239,
640
+ "step": 900
641
+ },
642
+ {
643
+ "epoch": 1.1084704448507008,
644
+ "grad_norm": 0.23486360907554626,
645
+ "learning_rate": 5e-05,
646
+ "loss": 0.7335,
647
+ "step": 910
648
+ },
649
+ {
650
+ "epoch": 1.1206581352833638,
651
+ "grad_norm": 0.22189787030220032,
652
+ "learning_rate": 5e-05,
653
+ "loss": 0.7149,
654
+ "step": 920
655
+ },
656
+ {
657
+ "epoch": 1.1328458257160268,
658
+ "grad_norm": 0.25818824768066406,
659
+ "learning_rate": 5e-05,
660
+ "loss": 0.7322,
661
+ "step": 930
662
+ },
663
+ {
664
+ "epoch": 1.1450335161486898,
665
+ "grad_norm": 0.2088775783777237,
666
+ "learning_rate": 5e-05,
667
+ "loss": 0.7382,
668
+ "step": 940
669
+ },
670
+ {
671
+ "epoch": 1.1572212065813527,
672
+ "grad_norm": 0.19445833563804626,
673
+ "learning_rate": 5e-05,
674
+ "loss": 0.7223,
675
+ "step": 950
676
+ },
677
+ {
678
+ "epoch": 1.1694088970140157,
679
+ "grad_norm": 0.21489456295967102,
680
+ "learning_rate": 5e-05,
681
+ "loss": 0.7317,
682
+ "step": 960
683
+ },
684
+ {
685
+ "epoch": 1.181596587446679,
686
+ "grad_norm": 0.24050794541835785,
687
+ "learning_rate": 5e-05,
688
+ "loss": 0.7275,
689
+ "step": 970
690
+ },
691
+ {
692
+ "epoch": 1.193784277879342,
693
+ "grad_norm": 0.18142102658748627,
694
+ "learning_rate": 5e-05,
695
+ "loss": 0.7046,
696
+ "step": 980
697
+ },
698
+ {
699
+ "epoch": 1.205971968312005,
700
+ "grad_norm": 0.4781077802181244,
701
+ "learning_rate": 5e-05,
702
+ "loss": 0.7433,
703
+ "step": 990
704
+ },
705
+ {
706
+ "epoch": 1.218159658744668,
707
+ "grad_norm": 0.20628023147583008,
708
+ "learning_rate": 5e-05,
709
+ "loss": 0.7403,
710
+ "step": 1000
711
+ },
712
+ {
713
+ "epoch": 1.230347349177331,
714
+ "grad_norm": 0.19117140769958496,
715
+ "learning_rate": 5e-05,
716
+ "loss": 0.7308,
717
+ "step": 1010
718
+ },
719
+ {
720
+ "epoch": 1.242535039609994,
721
+ "grad_norm": 0.2107745110988617,
722
+ "learning_rate": 5e-05,
723
+ "loss": 0.7201,
724
+ "step": 1020
725
+ },
726
+ {
727
+ "epoch": 1.2547227300426569,
728
+ "grad_norm": 0.20551154017448425,
729
+ "learning_rate": 5e-05,
730
+ "loss": 0.7038,
731
+ "step": 1030
732
+ },
733
+ {
734
+ "epoch": 1.2669104204753199,
735
+ "grad_norm": 0.26577022671699524,
736
+ "learning_rate": 5e-05,
737
+ "loss": 0.7096,
738
+ "step": 1040
739
+ },
740
+ {
741
+ "epoch": 1.2790981109079829,
742
+ "grad_norm": 0.22094209492206573,
743
+ "learning_rate": 5e-05,
744
+ "loss": 0.718,
745
+ "step": 1050
746
+ },
747
+ {
748
+ "epoch": 1.291285801340646,
749
+ "grad_norm": 0.18224076926708221,
750
+ "learning_rate": 5e-05,
751
+ "loss": 0.7261,
752
+ "step": 1060
753
+ },
754
+ {
755
+ "epoch": 1.303473491773309,
756
+ "grad_norm": 0.226251482963562,
757
+ "learning_rate": 5e-05,
758
+ "loss": 0.7372,
759
+ "step": 1070
760
+ },
761
+ {
762
+ "epoch": 1.315661182205972,
763
+ "grad_norm": 0.20186461508274078,
764
+ "learning_rate": 5e-05,
765
+ "loss": 0.7074,
766
+ "step": 1080
767
+ },
768
+ {
769
+ "epoch": 1.327848872638635,
770
+ "grad_norm": 0.22229762375354767,
771
+ "learning_rate": 5e-05,
772
+ "loss": 0.7069,
773
+ "step": 1090
774
+ },
775
+ {
776
+ "epoch": 1.340036563071298,
777
+ "grad_norm": 0.1776452213525772,
778
+ "learning_rate": 5e-05,
779
+ "loss": 0.714,
780
+ "step": 1100
781
+ },
782
+ {
783
+ "epoch": 1.352224253503961,
784
+ "grad_norm": 0.20680229365825653,
785
+ "learning_rate": 5e-05,
786
+ "loss": 0.7066,
787
+ "step": 1110
788
+ },
789
+ {
790
+ "epoch": 1.364411943936624,
791
+ "grad_norm": 0.21842031180858612,
792
+ "learning_rate": 5e-05,
793
+ "loss": 0.7281,
794
+ "step": 1120
795
+ },
796
+ {
797
+ "epoch": 1.376599634369287,
798
+ "grad_norm": 0.28037774562835693,
799
+ "learning_rate": 5e-05,
800
+ "loss": 0.7149,
801
+ "step": 1130
802
+ },
803
+ {
804
+ "epoch": 1.38878732480195,
805
+ "grad_norm": 0.20962920784950256,
806
+ "learning_rate": 5e-05,
807
+ "loss": 0.6951,
808
+ "step": 1140
809
+ },
810
+ {
811
+ "epoch": 1.400975015234613,
812
+ "grad_norm": 0.18751871585845947,
813
+ "learning_rate": 5e-05,
814
+ "loss": 0.7101,
815
+ "step": 1150
816
+ },
817
+ {
818
+ "epoch": 1.413162705667276,
819
+ "grad_norm": 0.2722964584827423,
820
+ "learning_rate": 5e-05,
821
+ "loss": 0.7205,
822
+ "step": 1160
823
+ },
824
+ {
825
+ "epoch": 1.425350396099939,
826
+ "grad_norm": 0.20932041108608246,
827
+ "learning_rate": 5e-05,
828
+ "loss": 0.7082,
829
+ "step": 1170
830
+ },
831
+ {
832
+ "epoch": 1.437538086532602,
833
+ "grad_norm": 0.17836257815361023,
834
+ "learning_rate": 5e-05,
835
+ "loss": 0.7232,
836
+ "step": 1180
837
+ },
838
+ {
839
+ "epoch": 1.449725776965265,
840
+ "grad_norm": 0.3180703818798065,
841
+ "learning_rate": 5e-05,
842
+ "loss": 0.7245,
843
+ "step": 1190
844
+ },
845
+ {
846
+ "epoch": 1.4619134673979282,
847
+ "grad_norm": 0.201424703001976,
848
+ "learning_rate": 5e-05,
849
+ "loss": 0.7232,
850
+ "step": 1200
851
+ },
852
+ {
853
+ "epoch": 1.4741011578305911,
854
+ "grad_norm": 0.20228999853134155,
855
+ "learning_rate": 5e-05,
856
+ "loss": 0.7169,
857
+ "step": 1210
858
+ },
859
+ {
860
+ "epoch": 1.4862888482632541,
861
+ "grad_norm": 0.17478394508361816,
862
+ "learning_rate": 5e-05,
863
+ "loss": 0.704,
864
+ "step": 1220
865
+ },
866
+ {
867
+ "epoch": 1.4984765386959171,
868
+ "grad_norm": 0.20061351358890533,
869
+ "learning_rate": 5e-05,
870
+ "loss": 0.6792,
871
+ "step": 1230
872
+ }
873
+ ],
874
+ "logging_steps": 10,
875
+ "max_steps": 2463,
876
+ "num_input_tokens_seen": 0,
877
+ "num_train_epochs": 3,
878
+ "save_steps": 205,
879
+ "stateful_callbacks": {
880
+ "TrainerControl": {
881
+ "args": {
882
+ "should_epoch_stop": false,
883
+ "should_evaluate": false,
884
+ "should_log": false,
885
+ "should_save": true,
886
+ "should_training_stop": false
887
+ },
888
+ "attributes": {}
889
+ }
890
+ },
891
+ "total_flos": 5.045988908773409e+18,
892
+ "train_batch_size": 4,
893
+ "trial_name": null,
894
+ "trial_params": null
895
+ }
openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1230/training_args.bin ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:56896ea21940e33d7e2934267dff2a5b4a4a80195a0f165245f6617a0f88b90b
3
+ size 7377
openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1230/zero_to_fp32.py ADDED
@@ -0,0 +1,760 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python
2
+
3
+ # Copyright (c) Microsoft Corporation.
4
+ # SPDX-License-Identifier: Apache-2.0
5
+
6
+ # DeepSpeed Team
7
+
8
+ # This script extracts fp32 consolidated weights from a zero 1, 2 and 3 DeepSpeed checkpoints. It gets
9
+ # copied into the top level checkpoint dir, so the user can easily do the conversion at any point in
10
+ # the future. Once extracted, the weights don't require DeepSpeed and can be used in any
11
+ # application.
12
+ #
13
+ # example:
14
+ # python zero_to_fp32.py . output_dir/
15
+ # or
16
+ # python zero_to_fp32.py . output_dir/ --safe_serialization
17
+
18
+ import argparse
19
+ import torch
20
+ import glob
21
+ import math
22
+ import os
23
+ import re
24
+ import gc
25
+ import json
26
+ import numpy as np
27
+ from tqdm import tqdm
28
+ from collections import OrderedDict
29
+ from dataclasses import dataclass
30
+
31
+ # while this script doesn't use deepspeed to recover data, since the checkpoints are pickled with
32
+ # DeepSpeed data structures it has to be available in the current python environment.
33
+ from deepspeed.utils import logger
34
+ from deepspeed.checkpoint.constants import (DS_VERSION, OPTIMIZER_STATE_DICT, SINGLE_PARTITION_OF_FP32_GROUPS,
35
+ FP32_FLAT_GROUPS, ZERO_STAGE, PARTITION_COUNT, PARAM_SHAPES, BUFFER_NAMES,
36
+ FROZEN_PARAM_SHAPES, FROZEN_PARAM_FRAGMENTS)
37
+
38
+
39
+ @dataclass
40
+ class zero_model_state:
41
+ buffers: dict()
42
+ param_shapes: dict()
43
+ shared_params: list
44
+ ds_version: int
45
+ frozen_param_shapes: dict()
46
+ frozen_param_fragments: dict()
47
+
48
+
49
+ debug = 0
50
+
51
+ # load to cpu
52
+ device = torch.device('cpu')
53
+
54
+
55
+ def atoi(text):
56
+ return int(text) if text.isdigit() else text
57
+
58
+
59
+ def natural_keys(text):
60
+ '''
61
+ alist.sort(key=natural_keys) sorts in human order
62
+ http://nedbatchelder.com/blog/200712/human_sorting.html
63
+ (See Toothy's implementation in the comments)
64
+ '''
65
+ return [atoi(c) for c in re.split(r'(\d+)', text)]
66
+
67
+
68
+ def get_model_state_file(checkpoint_dir, zero_stage):
69
+ if not os.path.isdir(checkpoint_dir):
70
+ raise FileNotFoundError(f"Directory '{checkpoint_dir}' doesn't exist")
71
+
72
+ # there should be only one file
73
+ if zero_stage <= 2:
74
+ file = os.path.join(checkpoint_dir, "mp_rank_00_model_states.pt")
75
+ elif zero_stage == 3:
76
+ file = os.path.join(checkpoint_dir, "zero_pp_rank_0_mp_rank_00_model_states.pt")
77
+
78
+ if not os.path.exists(file):
79
+ raise FileNotFoundError(f"can't find model states file at '{file}'")
80
+
81
+ return file
82
+
83
+
84
+ def get_checkpoint_files(checkpoint_dir, glob_pattern):
85
+ # XXX: need to test that this simple glob rule works for multi-node setup too
86
+ ckpt_files = sorted(glob.glob(os.path.join(checkpoint_dir, glob_pattern)), key=natural_keys)
87
+
88
+ if len(ckpt_files) == 0:
89
+ raise FileNotFoundError(f"can't find {glob_pattern} files in directory '{checkpoint_dir}'")
90
+
91
+ return ckpt_files
92
+
93
+
94
+ def get_optim_files(checkpoint_dir):
95
+ return get_checkpoint_files(checkpoint_dir, "*_optim_states.pt")
96
+
97
+
98
+ def get_model_state_files(checkpoint_dir):
99
+ return get_checkpoint_files(checkpoint_dir, "*_model_states.pt")
100
+
101
+
102
+ def parse_model_states(files):
103
+ zero_model_states = []
104
+ for file in files:
105
+ state_dict = torch.load(file, map_location=device, weights_only=False)
106
+
107
+ if BUFFER_NAMES not in state_dict:
108
+ raise ValueError(f"{file} is not a model state checkpoint")
109
+ buffer_names = state_dict[BUFFER_NAMES]
110
+ if debug:
111
+ print("Found buffers:", buffer_names)
112
+
113
+ # recover just the buffers while restoring them to fp32 if they were saved in fp16
114
+ buffers = {k: v.float() for k, v in state_dict["module"].items() if k in buffer_names}
115
+ param_shapes = state_dict[PARAM_SHAPES]
116
+
117
+ # collect parameters that are included in param_shapes
118
+ param_names = []
119
+ for s in param_shapes:
120
+ for name in s.keys():
121
+ param_names.append(name)
122
+
123
+ # update with frozen parameters
124
+ frozen_param_shapes = state_dict.get(FROZEN_PARAM_SHAPES, None)
125
+ if frozen_param_shapes is not None:
126
+ if debug:
127
+ print(f"Found frozen_param_shapes: {frozen_param_shapes}")
128
+ param_names += list(frozen_param_shapes.keys())
129
+
130
+ # handle shared params
131
+ shared_params = [[k, v] for k, v in state_dict["shared_params"].items()]
132
+
133
+ ds_version = state_dict.get(DS_VERSION, None)
134
+
135
+ frozen_param_fragments = state_dict.get(FROZEN_PARAM_FRAGMENTS, None)
136
+
137
+ z_model_state = zero_model_state(buffers=buffers,
138
+ param_shapes=param_shapes,
139
+ shared_params=shared_params,
140
+ ds_version=ds_version,
141
+ frozen_param_shapes=frozen_param_shapes,
142
+ frozen_param_fragments=frozen_param_fragments)
143
+ zero_model_states.append(z_model_state)
144
+
145
+ return zero_model_states
146
+
147
+
148
+ def parse_optim_states(files, ds_checkpoint_dir):
149
+ total_files = len(files)
150
+ state_dicts = []
151
+ for f in tqdm(files, desc='Loading checkpoint shards'):
152
+ state_dict = torch.load(f, map_location=device, mmap=True, weights_only=False)
153
+ # immediately discard the potentially huge 2 optimizer states as we only care for fp32 master weights
154
+ # and also handle the case where it was already removed by another helper script
155
+ state_dict["optimizer_state_dict"].pop("optimizer_state_dict", None)
156
+ state_dicts.append(state_dict)
157
+
158
+ if ZERO_STAGE not in state_dicts[0][OPTIMIZER_STATE_DICT]:
159
+ raise ValueError(f"{files[0]} is not a zero checkpoint")
160
+ zero_stage = state_dicts[0][OPTIMIZER_STATE_DICT][ZERO_STAGE]
161
+ world_size = state_dicts[0][OPTIMIZER_STATE_DICT][PARTITION_COUNT]
162
+
163
+ # For ZeRO-2 each param group can have different partition_count as data parallelism for expert
164
+ # parameters can be different from data parallelism for non-expert parameters. So we can just
165
+ # use the max of the partition_count to get the dp world_size.
166
+
167
+ if type(world_size) is list:
168
+ world_size = max(world_size)
169
+
170
+ if world_size != total_files:
171
+ raise ValueError(
172
+ f"Expected {world_size} of '*_optim_states.pt' under '{ds_checkpoint_dir}' but found {total_files} files. "
173
+ "Possibly due to an overwrite of an old checkpoint, or a checkpoint didn't get saved by one or more processes."
174
+ )
175
+
176
+ # the groups are named differently in each stage
177
+ if zero_stage <= 2:
178
+ fp32_groups_key = SINGLE_PARTITION_OF_FP32_GROUPS
179
+ elif zero_stage == 3:
180
+ fp32_groups_key = FP32_FLAT_GROUPS
181
+ else:
182
+ raise ValueError(f"unknown zero stage {zero_stage}")
183
+
184
+ fp32_flat_groups = [state_dicts[i][OPTIMIZER_STATE_DICT][fp32_groups_key] for i in range(len(state_dicts))]
185
+ return zero_stage, world_size, fp32_flat_groups
186
+
187
+
188
+ def _get_fp32_state_dict_from_zero_checkpoint(ds_checkpoint_dir, exclude_frozen_parameters):
189
+ """
190
+ Returns fp32 state_dict reconstructed from ds checkpoint
191
+
192
+ Args:
193
+ - ``ds_checkpoint_dir``: path to the deepspeed checkpoint folder (where the optimizer files are)
194
+
195
+ """
196
+ print(f"Processing zero checkpoint '{ds_checkpoint_dir}'")
197
+
198
+ optim_files = get_optim_files(ds_checkpoint_dir)
199
+ zero_stage, world_size, fp32_flat_groups = parse_optim_states(optim_files, ds_checkpoint_dir)
200
+ print(f"Detected checkpoint of type zero stage {zero_stage}, world_size: {world_size}")
201
+
202
+ model_files = get_model_state_files(ds_checkpoint_dir)
203
+
204
+ zero_model_states = parse_model_states(model_files)
205
+ print(f'Parsing checkpoint created by deepspeed=={zero_model_states[0].ds_version}')
206
+
207
+ if zero_stage <= 2:
208
+ return _get_fp32_state_dict_from_zero2_checkpoint(world_size, fp32_flat_groups, zero_model_states,
209
+ exclude_frozen_parameters)
210
+ elif zero_stage == 3:
211
+ return _get_fp32_state_dict_from_zero3_checkpoint(world_size, fp32_flat_groups, zero_model_states,
212
+ exclude_frozen_parameters)
213
+
214
+
215
+ def _zero2_merge_frozen_params(state_dict, zero_model_states):
216
+ if zero_model_states[0].frozen_param_shapes is None or len(zero_model_states[0].frozen_param_shapes) == 0:
217
+ return
218
+
219
+ frozen_param_shapes = zero_model_states[0].frozen_param_shapes
220
+ frozen_param_fragments = zero_model_states[0].frozen_param_fragments
221
+
222
+ if debug:
223
+ num_elem = sum(s.numel() for s in frozen_param_shapes.values())
224
+ print(f'rank 0: {FROZEN_PARAM_SHAPES}.numel = {num_elem}')
225
+
226
+ wanted_params = len(frozen_param_shapes)
227
+ wanted_numel = sum(s.numel() for s in frozen_param_shapes.values())
228
+ avail_numel = sum([p.numel() for p in frozen_param_fragments.values()])
229
+ print(f'Frozen params: Have {avail_numel} numels to process.')
230
+ print(f'Frozen params: Need {wanted_numel} numels in {wanted_params} params')
231
+
232
+ total_params = 0
233
+ total_numel = 0
234
+ for name, shape in frozen_param_shapes.items():
235
+ total_params += 1
236
+ unpartitioned_numel = shape.numel()
237
+ total_numel += unpartitioned_numel
238
+
239
+ state_dict[name] = frozen_param_fragments[name]
240
+
241
+ if debug:
242
+ print(f"{name} full shape: {shape} unpartitioned numel {unpartitioned_numel} ")
243
+
244
+ print(f"Reconstructed Frozen fp32 state dict with {total_params} params {total_numel} elements")
245
+
246
+
247
+ def _has_callable(obj, fn):
248
+ attr = getattr(obj, fn, None)
249
+ return callable(attr)
250
+
251
+
252
+ def _zero2_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states):
253
+ param_shapes = zero_model_states[0].param_shapes
254
+
255
+ # Reconstruction protocol:
256
+ #
257
+ # XXX: document this
258
+
259
+ if debug:
260
+ for i in range(world_size):
261
+ for j in range(len(fp32_flat_groups[0])):
262
+ print(f"{FP32_FLAT_GROUPS}[{i}][{j}].shape={fp32_flat_groups[i][j].shape}")
263
+
264
+ # XXX: memory usage doubles here (zero2)
265
+ num_param_groups = len(fp32_flat_groups[0])
266
+ merged_single_partition_of_fp32_groups = []
267
+ for i in range(num_param_groups):
268
+ merged_partitions = [sd[i] for sd in fp32_flat_groups]
269
+ full_single_fp32_vector = torch.cat(merged_partitions, 0)
270
+ merged_single_partition_of_fp32_groups.append(full_single_fp32_vector)
271
+ avail_numel = sum(
272
+ [full_single_fp32_vector.numel() for full_single_fp32_vector in merged_single_partition_of_fp32_groups])
273
+
274
+ if debug:
275
+ wanted_params = sum([len(shapes) for shapes in param_shapes])
276
+ wanted_numel = sum([sum(shape.numel() for shape in shapes.values()) for shapes in param_shapes])
277
+ # not asserting if there is a mismatch due to possible padding
278
+ print(f"Have {avail_numel} numels to process.")
279
+ print(f"Need {wanted_numel} numels in {wanted_params} params.")
280
+
281
+ # params
282
+ # XXX: for huge models that can't fit into the host's RAM we will have to recode this to support
283
+ # out-of-core computing solution
284
+ total_numel = 0
285
+ total_params = 0
286
+ for shapes, full_single_fp32_vector in zip(param_shapes, merged_single_partition_of_fp32_groups):
287
+ offset = 0
288
+ avail_numel = full_single_fp32_vector.numel()
289
+ for name, shape in shapes.items():
290
+
291
+ unpartitioned_numel = shape.numel() if _has_callable(shape, 'numel') else math.prod(shape)
292
+ total_numel += unpartitioned_numel
293
+ total_params += 1
294
+
295
+ if debug:
296
+ print(f"{name} full shape: {shape} unpartitioned numel {unpartitioned_numel} ")
297
+ state_dict[name] = full_single_fp32_vector.narrow(0, offset, unpartitioned_numel).view(shape)
298
+ offset += unpartitioned_numel
299
+
300
+ # Z2 started to align to 2*world_size to improve nccl performance. Therefore both offset and
301
+ # avail_numel can differ by anywhere between 0..2*world_size. Due to two unrelated complex
302
+ # paddings performed in the code it's almost impossible to predict the exact numbers w/o the
303
+ # live optimizer object, so we are checking that the numbers are within the right range
304
+ align_to = 2 * world_size
305
+
306
+ def zero2_align(x):
307
+ return align_to * math.ceil(x / align_to)
308
+
309
+ if debug:
310
+ print(f"original offset={offset}, avail_numel={avail_numel}")
311
+
312
+ offset = zero2_align(offset)
313
+ avail_numel = zero2_align(avail_numel)
314
+
315
+ if debug:
316
+ print(f"aligned offset={offset}, avail_numel={avail_numel}")
317
+
318
+ # Sanity check
319
+ if offset != avail_numel:
320
+ raise ValueError(f"consumed {offset} numels out of {avail_numel} - something is wrong")
321
+
322
+ print(f"Reconstructed fp32 state dict with {total_params} params {total_numel} elements")
323
+
324
+
325
+ def _get_fp32_state_dict_from_zero2_checkpoint(world_size, fp32_flat_groups, zero_model_states,
326
+ exclude_frozen_parameters):
327
+ state_dict = OrderedDict()
328
+
329
+ # buffers
330
+ buffers = zero_model_states[0].buffers
331
+ state_dict.update(buffers)
332
+ if debug:
333
+ print(f"added {len(buffers)} buffers")
334
+
335
+ if not exclude_frozen_parameters:
336
+ _zero2_merge_frozen_params(state_dict, zero_model_states)
337
+
338
+ _zero2_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states)
339
+
340
+ # recover shared parameters
341
+ for pair in zero_model_states[0].shared_params:
342
+ if pair[1] in state_dict:
343
+ state_dict[pair[0]] = state_dict[pair[1]]
344
+
345
+ return state_dict
346
+
347
+
348
+ def zero3_partitioned_param_info(unpartitioned_numel, world_size):
349
+ remainder = unpartitioned_numel % world_size
350
+ padding_numel = (world_size - remainder) if remainder else 0
351
+ partitioned_numel = math.ceil(unpartitioned_numel / world_size)
352
+ return partitioned_numel, padding_numel
353
+
354
+
355
+ def _zero3_merge_frozen_params(state_dict, world_size, zero_model_states):
356
+ if zero_model_states[0].frozen_param_shapes is None or len(zero_model_states[0].frozen_param_shapes) == 0:
357
+ return
358
+
359
+ if debug:
360
+ for i in range(world_size):
361
+ num_elem = sum(s.numel() for s in zero_model_states[i].frozen_param_fragments.values())
362
+ print(f'rank {i}: {FROZEN_PARAM_SHAPES}.numel = {num_elem}')
363
+
364
+ frozen_param_shapes = zero_model_states[0].frozen_param_shapes
365
+ wanted_params = len(frozen_param_shapes)
366
+ wanted_numel = sum(s.numel() for s in frozen_param_shapes.values())
367
+ avail_numel = sum([p.numel() for p in zero_model_states[0].frozen_param_fragments.values()]) * world_size
368
+ print(f'Frozen params: Have {avail_numel} numels to process.')
369
+ print(f'Frozen params: Need {wanted_numel} numels in {wanted_params} params')
370
+
371
+ total_params = 0
372
+ total_numel = 0
373
+ for name, shape in zero_model_states[0].frozen_param_shapes.items():
374
+ total_params += 1
375
+ unpartitioned_numel = shape.numel()
376
+ total_numel += unpartitioned_numel
377
+
378
+ param_frags = tuple(model_state.frozen_param_fragments[name] for model_state in zero_model_states)
379
+ state_dict[name] = torch.cat(param_frags, 0).narrow(0, 0, unpartitioned_numel).view(shape)
380
+
381
+ partitioned_numel, partitioned_padding_numel = zero3_partitioned_param_info(unpartitioned_numel, world_size)
382
+
383
+ if debug:
384
+ print(
385
+ f"Frozen params: {total_params} {name} full shape: {shape} partition0 numel={partitioned_numel} partitioned_padding_numel={partitioned_padding_numel}"
386
+ )
387
+
388
+ print(f"Reconstructed Frozen fp32 state dict with {total_params} params {total_numel} elements")
389
+
390
+
391
+ class GatheredTensor:
392
+ """
393
+ A pseudo tensor that collects partitioned weights.
394
+ It is more memory efficient when there are multiple groups.
395
+ """
396
+
397
+ def __init__(self, flat_groups, flat_groups_offset, offset, partitioned_numel, shape):
398
+ self.flat_groups = flat_groups
399
+ self.flat_groups_offset = flat_groups_offset
400
+ self.offset = offset
401
+ self.partitioned_numel = partitioned_numel
402
+ self.shape = shape
403
+ self.dtype = self.flat_groups[0][0].dtype
404
+
405
+ def contiguous(self):
406
+ """
407
+ Merge partitioned weights from flat_groups into a single tensor.
408
+ """
409
+ end_idx = self.offset + self.partitioned_numel
410
+ world_size = len(self.flat_groups)
411
+ pad_flat_param_chunks = []
412
+
413
+ for rank_i in range(world_size):
414
+ # for each rank, we need to collect weights from related group/groups
415
+ flat_groups_at_rank_i = self.flat_groups[rank_i]
416
+ start_group_id = None
417
+ end_group_id = None
418
+ for group_id in range(len(self.flat_groups_offset)):
419
+ if self.flat_groups_offset[group_id] <= self.offset < self.flat_groups_offset[group_id + 1]:
420
+ start_group_id = group_id
421
+ if self.flat_groups_offset[group_id] < end_idx <= self.flat_groups_offset[group_id + 1]:
422
+ end_group_id = group_id
423
+ break
424
+ # collect weights from related group/groups
425
+ for group_id in range(start_group_id, end_group_id + 1):
426
+ flat_tensor = flat_groups_at_rank_i[group_id]
427
+ start_offset = self.offset - self.flat_groups_offset[group_id]
428
+ end_offset = min(end_idx, self.flat_groups_offset[group_id + 1]) - self.flat_groups_offset[group_id]
429
+ pad_flat_param_chunks.append(flat_tensor[start_offset:end_offset])
430
+
431
+ # collect weights from all ranks
432
+ pad_flat_param = torch.cat(pad_flat_param_chunks, dim=0)
433
+ param = pad_flat_param[:self.shape.numel()].view(self.shape).contiguous()
434
+ return param
435
+
436
+
437
+ def _zero3_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states):
438
+ param_shapes = zero_model_states[0].param_shapes
439
+ avail_numel = sum([flat_group.numel() for flat_group in fp32_flat_groups[0]]) * world_size
440
+
441
+ # Reconstruction protocol: For zero3 we need to zip the partitions together at boundary of each
442
+ # param, re-consolidating each param, while dealing with padding if any
443
+
444
+ # merge list of dicts, preserving order
445
+ param_shapes = {k: v for d in param_shapes for k, v in d.items()}
446
+
447
+ if debug:
448
+ for i in range(world_size):
449
+ print(f"{FP32_FLAT_GROUPS}[{i}].shape={fp32_flat_groups[i].shape}")
450
+
451
+ wanted_params = len(param_shapes)
452
+ wanted_numel = sum(shape.numel() for shape in param_shapes.values())
453
+ # not asserting if there is a mismatch due to possible padding
454
+ avail_numel = fp32_flat_groups[0].numel() * world_size
455
+ print(f"Trainable params: Have {avail_numel} numels to process.")
456
+ print(f"Trainable params: Need {wanted_numel} numels in {wanted_params} params.")
457
+
458
+ # params
459
+ # XXX: for huge models that can't fit into the host's RAM we will have to recode this to support
460
+ # out-of-core computing solution
461
+ offset = 0
462
+ total_numel = 0
463
+ total_params = 0
464
+ flat_groups_offset = [0] + list(np.cumsum([flat_tensor.numel() for flat_tensor in fp32_flat_groups[0]]))
465
+ for name, shape in tqdm(param_shapes.items(), desc='Gathering sharded weights'):
466
+ unpartitioned_numel = shape.numel()
467
+ total_numel += unpartitioned_numel
468
+ total_params += 1
469
+ partitioned_numel, partitioned_padding_numel = zero3_partitioned_param_info(unpartitioned_numel, world_size)
470
+
471
+ if debug:
472
+ print(
473
+ f"Trainable params: {total_params} {name} full shape: {shape} partition0 numel={partitioned_numel} partitioned_padding_numel={partitioned_padding_numel}"
474
+ )
475
+
476
+ # memory efficient tensor
477
+ tensor = GatheredTensor(fp32_flat_groups, flat_groups_offset, offset, partitioned_numel, shape)
478
+ state_dict[name] = tensor
479
+ offset += partitioned_numel
480
+
481
+ offset *= world_size
482
+
483
+ # Sanity check
484
+ if offset != avail_numel:
485
+ raise ValueError(f"consumed {offset} numels out of {avail_numel} - something is wrong")
486
+
487
+ print(f"Reconstructed Trainable fp32 state dict with {total_params} params {total_numel} elements")
488
+
489
+
490
+ def _get_fp32_state_dict_from_zero3_checkpoint(world_size, fp32_flat_groups, zero_model_states,
491
+ exclude_frozen_parameters):
492
+ state_dict = OrderedDict()
493
+
494
+ # buffers
495
+ buffers = zero_model_states[0].buffers
496
+ state_dict.update(buffers)
497
+ if debug:
498
+ print(f"added {len(buffers)} buffers")
499
+
500
+ if not exclude_frozen_parameters:
501
+ _zero3_merge_frozen_params(state_dict, world_size, zero_model_states)
502
+
503
+ _zero3_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states)
504
+
505
+ # recover shared parameters
506
+ for pair in zero_model_states[0].shared_params:
507
+ if pair[1] in state_dict:
508
+ state_dict[pair[0]] = state_dict[pair[1]]
509
+
510
+ return state_dict
511
+
512
+
513
+ def to_torch_tensor(state_dict, return_empty_tensor=False):
514
+ """
515
+ Convert state_dict of GatheredTensor to torch tensor
516
+ """
517
+ torch_state_dict = {}
518
+ converted_tensors = {}
519
+ for name, tensor in state_dict.items():
520
+ tensor_id = id(tensor)
521
+ if tensor_id in converted_tensors: # shared tensors
522
+ shared_tensor = torch_state_dict[converted_tensors[tensor_id]]
523
+ torch_state_dict[name] = shared_tensor
524
+ else:
525
+ converted_tensors[tensor_id] = name
526
+ if return_empty_tensor:
527
+ torch_state_dict[name] = torch.empty(tensor.shape, dtype=tensor.dtype)
528
+ else:
529
+ torch_state_dict[name] = tensor.contiguous()
530
+ return torch_state_dict
531
+
532
+
533
+ def get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir,
534
+ tag=None,
535
+ exclude_frozen_parameters=False,
536
+ lazy_mode=False):
537
+ """
538
+ Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated state_dict that can be loaded with
539
+ ``load_state_dict()`` and used for training without DeepSpeed or shared with others, for example
540
+ via a model hub.
541
+
542
+ Args:
543
+ - ``checkpoint_dir``: path to the desired checkpoint folder
544
+ - ``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``
545
+ - ``exclude_frozen_parameters``: exclude frozen parameters
546
+ - ``lazy_mode``: get state_dict in lazy mode. It returns a dict of pesduo tensor instead of torch tensor, which is more memory efficient.
547
+ Convert the pesduo tensor to torch tensor by ``.contiguous()``
548
+
549
+ Returns:
550
+ - pytorch ``state_dict``
551
+
552
+ A typical usage might be ::
553
+
554
+ from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint
555
+ # do the training and checkpoint saving
556
+ state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir) # already on cpu
557
+ model = model.cpu() # move to cpu
558
+ model.load_state_dict(state_dict)
559
+ # submit to model hub or save the model to share with others
560
+
561
+ In this example the ``model`` will no longer be usable in the deepspeed context of the same
562
+ application. i.e. you will need to re-initialize the deepspeed engine, since
563
+ ``model.load_state_dict(state_dict)`` will remove all the deepspeed magic from it.
564
+
565
+ If you want it all done for you, use ``load_state_dict_from_zero_checkpoint`` instead.
566
+
567
+ Note: the above usage may not work if your application doesn't have sufficient free CPU memory.
568
+ You may need to use the offline approach using the ``zero_to_fp32.py`` script that is saved with
569
+ the checkpoint. Or you can load state_dict in lazy mode ::
570
+
571
+ from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint
572
+ state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, lazy_mode=True) # not on cpu
573
+ for name, lazy_tensor in state_dict.item():
574
+ tensor = lazy_tensor.contiguous() # to cpu
575
+ print(name, tensor)
576
+ # del tensor to release memory if it no longer in use
577
+ """
578
+ if tag is None:
579
+ latest_path = os.path.join(checkpoint_dir, 'latest')
580
+ if os.path.isfile(latest_path):
581
+ with open(latest_path, 'r') as fd:
582
+ tag = fd.read().strip()
583
+ else:
584
+ raise ValueError(f"Unable to find 'latest' file at {latest_path}")
585
+
586
+ ds_checkpoint_dir = os.path.join(checkpoint_dir, tag)
587
+
588
+ if not os.path.isdir(ds_checkpoint_dir):
589
+ raise FileNotFoundError(f"Directory '{ds_checkpoint_dir}' doesn't exist")
590
+
591
+ state_dict = _get_fp32_state_dict_from_zero_checkpoint(ds_checkpoint_dir, exclude_frozen_parameters)
592
+ if lazy_mode:
593
+ return state_dict
594
+ else:
595
+ return to_torch_tensor(state_dict)
596
+
597
+
598
+ def convert_zero_checkpoint_to_fp32_state_dict(checkpoint_dir,
599
+ output_dir,
600
+ max_shard_size="5GB",
601
+ safe_serialization=False,
602
+ tag=None,
603
+ exclude_frozen_parameters=False):
604
+ """
605
+ Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated ``state_dict`` file that can be
606
+ loaded with ``torch.load(file)`` + ``load_state_dict()`` and used for training without DeepSpeed.
607
+
608
+ Args:
609
+ - ``checkpoint_dir``: path to the desired checkpoint folder. (one that contains the tag-folder, like ``global_step14``)
610
+ - ``output_dir``: directory to the pytorch fp32 state_dict output files
611
+ - ``max_shard_size``: the maximum size for a checkpoint before being sharded, default value is 5GB
612
+ - ``safe_serialization``: whether to save the model using `safetensors` or the traditional PyTorch way (that uses `pickle`).
613
+ - ``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``
614
+ - ``exclude_frozen_parameters``: exclude frozen parameters
615
+ """
616
+
617
+ # Dependency pre-check
618
+ if safe_serialization:
619
+ try:
620
+ from safetensors.torch import save_file
621
+ except ImportError:
622
+ print('If you want to use `safe_serialization`, please `pip install safetensors`')
623
+ raise
624
+ if max_shard_size is not None:
625
+ try:
626
+ from huggingface_hub import split_torch_state_dict_into_shards
627
+ except ImportError:
628
+ print('If you want to use `max_shard_size`, please `pip install huggingface_hub`')
629
+ raise
630
+
631
+ # Convert zero checkpoint to state_dict
632
+ state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir,
633
+ tag,
634
+ exclude_frozen_parameters,
635
+ lazy_mode=True)
636
+
637
+ # Shard the model if it is too big.
638
+ weights_name = "model.safetensors" if safe_serialization else "pytorch_model.bin"
639
+ if max_shard_size is not None:
640
+ filename_pattern = weights_name.replace(".bin", "{suffix}.bin").replace(".safetensors", "{suffix}.safetensors")
641
+ # an memory-efficient approach for sharding
642
+ empty_state_dict = to_torch_tensor(state_dict, return_empty_tensor=True)
643
+ state_dict_split = split_torch_state_dict_into_shards(empty_state_dict,
644
+ filename_pattern=filename_pattern,
645
+ max_shard_size=max_shard_size)
646
+ else:
647
+ from collections import namedtuple
648
+ StateDictSplit = namedtuple("StateDictSplit", ["is_sharded", "filename_to_tensors"])
649
+ state_dict_split = StateDictSplit(is_sharded=False,
650
+ filename_to_tensors={weights_name: list(state_dict.keys())})
651
+
652
+ # Save the model by shard
653
+ os.makedirs(output_dir, exist_ok=True)
654
+ filename_to_tensors = state_dict_split.filename_to_tensors.items()
655
+ for shard_file, tensors in tqdm(filename_to_tensors, desc="Saving checkpoint shards"):
656
+ shard_state_dict = {tensor_name: state_dict[tensor_name] for tensor_name in tensors}
657
+ shard_state_dict = to_torch_tensor(shard_state_dict)
658
+ output_path = os.path.join(output_dir, shard_file)
659
+ if safe_serialization:
660
+ save_file(shard_state_dict, output_path, metadata={"format": "pt"})
661
+ else:
662
+ torch.save(shard_state_dict, output_path)
663
+ # release the memory of current shard
664
+ for tensor_name in list(shard_state_dict.keys()):
665
+ del state_dict[tensor_name]
666
+ del shard_state_dict[tensor_name]
667
+ del shard_state_dict
668
+ gc.collect()
669
+
670
+ # Save index if sharded
671
+ if state_dict_split.is_sharded:
672
+ index = {
673
+ "metadata": state_dict_split.metadata,
674
+ "weight_map": state_dict_split.tensor_to_filename,
675
+ }
676
+ save_index_file = "model.safetensors.index.json" if safe_serialization else "pytorch_model.bin.index.json"
677
+ save_index_file = os.path.join(output_dir, save_index_file)
678
+ with open(save_index_file, "w", encoding="utf-8") as f:
679
+ content = json.dumps(index, indent=2, sort_keys=True) + "\n"
680
+ f.write(content)
681
+
682
+
683
+ def load_state_dict_from_zero_checkpoint(model, checkpoint_dir, tag=None):
684
+ """
685
+ 1. Put the provided model to cpu
686
+ 2. Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated ``state_dict``
687
+ 3. Load it into the provided model
688
+
689
+ Args:
690
+ - ``model``: the model object to update
691
+ - ``checkpoint_dir``: path to the desired checkpoint folder. (one that contains the tag-folder, like ``global_step14``)
692
+ - ``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``
693
+
694
+ Returns:
695
+ - ``model`: modified model
696
+
697
+ Make sure you have plenty of CPU memory available before you call this function. If you don't
698
+ have enough use the ``zero_to_fp32.py`` utility to do the conversion. You will find it
699
+ conveniently placed for you in the checkpoint folder.
700
+
701
+ A typical usage might be ::
702
+
703
+ from deepspeed.utils.zero_to_fp32 import load_state_dict_from_zero_checkpoint
704
+ model = load_state_dict_from_zero_checkpoint(trainer.model, checkpoint_dir)
705
+ # submit to model hub or save the model to share with others
706
+
707
+ Note, that once this was run, the ``model`` will no longer be usable in the deepspeed context
708
+ of the same application. i.e. you will need to re-initialize the deepspeed engine, since
709
+ ``model.load_state_dict(state_dict)`` will remove all the deepspeed magic from it.
710
+
711
+ """
712
+ logger.info("Extracting fp32 weights")
713
+ state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, tag)
714
+
715
+ logger.info("Overwriting model with fp32 weights")
716
+ model = model.cpu()
717
+ model.load_state_dict(state_dict, strict=False)
718
+
719
+ return model
720
+
721
+
722
+ if __name__ == "__main__":
723
+ parser = argparse.ArgumentParser()
724
+ parser.add_argument("checkpoint_dir",
725
+ type=str,
726
+ help="path to the desired checkpoint folder, e.g., path/checkpoint-12")
727
+ parser.add_argument("output_dir",
728
+ type=str,
729
+ help="directory to the pytorch fp32 state_dict output files"
730
+ "(e.g. path/checkpoint-12-output/)")
731
+ parser.add_argument(
732
+ "--max_shard_size",
733
+ type=str,
734
+ default="5GB",
735
+ help="The maximum size for a checkpoint before being sharded. Checkpoints shard will then be each of size"
736
+ "lower than this size. If expressed as a string, needs to be digits followed by a unit (like `5MB`"
737
+ "We default it to 5GB in order for models to be able to run easily on free-tier google colab instances"
738
+ "without CPU OOM issues.")
739
+ parser.add_argument(
740
+ "--safe_serialization",
741
+ default=False,
742
+ action='store_true',
743
+ help="Whether to save the model using `safetensors` or the traditional PyTorch way (that uses `pickle`).")
744
+ parser.add_argument("-t",
745
+ "--tag",
746
+ type=str,
747
+ default=None,
748
+ help="checkpoint tag used as a unique identifier for checkpoint. e.g., global_step1")
749
+ parser.add_argument("--exclude_frozen_parameters", action='store_true', help="exclude frozen parameters")
750
+ parser.add_argument("-d", "--debug", action='store_true', help="enable debug")
751
+ args = parser.parse_args()
752
+
753
+ debug = args.debug
754
+
755
+ convert_zero_checkpoint_to_fp32_state_dict(args.checkpoint_dir,
756
+ args.output_dir,
757
+ max_shard_size=args.max_shard_size,
758
+ safe_serialization=args.safe_serialization,
759
+ tag=args.tag,
760
+ exclude_frozen_parameters=args.exclude_frozen_parameters)
openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1435/config.json ADDED
@@ -0,0 +1,39 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "architectures": [
3
+ "LlamaForCausalLM"
4
+ ],
5
+ "attention_bias": false,
6
+ "attention_dropout": 0.0,
7
+ "bos_token_id": 128000,
8
+ "dtype": "bfloat16",
9
+ "eos_token_id": [
10
+ 128001,
11
+ 128008,
12
+ 128009
13
+ ],
14
+ "head_dim": 64,
15
+ "hidden_act": "silu",
16
+ "hidden_size": 2048,
17
+ "initializer_range": 0.02,
18
+ "intermediate_size": 8192,
19
+ "max_position_embeddings": 131072,
20
+ "mlp_bias": false,
21
+ "model_type": "llama",
22
+ "num_attention_heads": 32,
23
+ "num_hidden_layers": 16,
24
+ "num_key_value_heads": 8,
25
+ "pretraining_tp": 1,
26
+ "rms_norm_eps": 1e-05,
27
+ "rope_scaling": {
28
+ "factor": 32.0,
29
+ "high_freq_factor": 4.0,
30
+ "low_freq_factor": 1.0,
31
+ "original_max_position_embeddings": 8192,
32
+ "rope_type": "llama3"
33
+ },
34
+ "rope_theta": 500000.0,
35
+ "tie_word_embeddings": true,
36
+ "transformers_version": "4.57.1",
37
+ "use_cache": false,
38
+ "vocab_size": 128256
39
+ }
openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1435/generation_config.json ADDED
@@ -0,0 +1,12 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "bos_token_id": 128000,
3
+ "do_sample": true,
4
+ "eos_token_id": [
5
+ 128001,
6
+ 128008,
7
+ 128009
8
+ ],
9
+ "temperature": 0.6,
10
+ "top_p": 0.9,
11
+ "transformers_version": "4.57.1"
12
+ }
openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1435/latest ADDED
@@ -0,0 +1 @@
 
 
1
+ global_step1435
openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1435/model.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:cbe31e0f0b3fbf43f52fe7b150aee5c28e90a332cb8459fea906932f5e368c93
3
+ size 2996982344
openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1435/scheduler.pt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:4e5c4ae215497f7a140302ca16ad638046b865f4344b7be03b9098b7c29a892c
3
+ size 1465
openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1435/trainer_state.json ADDED
@@ -0,0 +1,1035 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "best_global_step": null,
3
+ "best_metric": null,
4
+ "best_model_checkpoint": null,
5
+ "epoch": 1.7483241925655089,
6
+ "eval_steps": 500,
7
+ "global_step": 1435,
8
+ "is_hyper_param_search": false,
9
+ "is_local_process_zero": true,
10
+ "is_world_process_zero": true,
11
+ "log_history": [
12
+ {
13
+ "epoch": 0.01218769043266301,
14
+ "grad_norm": 0.36486390233039856,
15
+ "learning_rate": 5e-05,
16
+ "loss": 1.2213,
17
+ "step": 10
18
+ },
19
+ {
20
+ "epoch": 0.02437538086532602,
21
+ "grad_norm": 0.2953564524650574,
22
+ "learning_rate": 5e-05,
23
+ "loss": 1.0223,
24
+ "step": 20
25
+ },
26
+ {
27
+ "epoch": 0.03656307129798903,
28
+ "grad_norm": 0.2766735851764679,
29
+ "learning_rate": 5e-05,
30
+ "loss": 0.9589,
31
+ "step": 30
32
+ },
33
+ {
34
+ "epoch": 0.04875076173065204,
35
+ "grad_norm": 0.32770997285842896,
36
+ "learning_rate": 5e-05,
37
+ "loss": 0.9081,
38
+ "step": 40
39
+ },
40
+ {
41
+ "epoch": 0.06093845216331505,
42
+ "grad_norm": 0.3052276372909546,
43
+ "learning_rate": 5e-05,
44
+ "loss": 0.8752,
45
+ "step": 50
46
+ },
47
+ {
48
+ "epoch": 0.07312614259597806,
49
+ "grad_norm": 0.240146666765213,
50
+ "learning_rate": 5e-05,
51
+ "loss": 0.873,
52
+ "step": 60
53
+ },
54
+ {
55
+ "epoch": 0.08531383302864107,
56
+ "grad_norm": 0.23872187733650208,
57
+ "learning_rate": 5e-05,
58
+ "loss": 0.8587,
59
+ "step": 70
60
+ },
61
+ {
62
+ "epoch": 0.09750152346130408,
63
+ "grad_norm": 0.32638487219810486,
64
+ "learning_rate": 5e-05,
65
+ "loss": 0.8656,
66
+ "step": 80
67
+ },
68
+ {
69
+ "epoch": 0.10968921389396709,
70
+ "grad_norm": 0.25019028782844543,
71
+ "learning_rate": 5e-05,
72
+ "loss": 0.8685,
73
+ "step": 90
74
+ },
75
+ {
76
+ "epoch": 0.1218769043266301,
77
+ "grad_norm": 0.2507675886154175,
78
+ "learning_rate": 5e-05,
79
+ "loss": 0.8662,
80
+ "step": 100
81
+ },
82
+ {
83
+ "epoch": 0.1340645947592931,
84
+ "grad_norm": 0.27966755628585815,
85
+ "learning_rate": 5e-05,
86
+ "loss": 0.8428,
87
+ "step": 110
88
+ },
89
+ {
90
+ "epoch": 0.14625228519195613,
91
+ "grad_norm": 0.2622673213481903,
92
+ "learning_rate": 5e-05,
93
+ "loss": 0.8416,
94
+ "step": 120
95
+ },
96
+ {
97
+ "epoch": 0.15843997562461914,
98
+ "grad_norm": 0.3003966510295868,
99
+ "learning_rate": 5e-05,
100
+ "loss": 0.8528,
101
+ "step": 130
102
+ },
103
+ {
104
+ "epoch": 0.17062766605728213,
105
+ "grad_norm": 0.32308319211006165,
106
+ "learning_rate": 5e-05,
107
+ "loss": 0.8357,
108
+ "step": 140
109
+ },
110
+ {
111
+ "epoch": 0.18281535648994515,
112
+ "grad_norm": 0.27279096841812134,
113
+ "learning_rate": 5e-05,
114
+ "loss": 0.8432,
115
+ "step": 150
116
+ },
117
+ {
118
+ "epoch": 0.19500304692260817,
119
+ "grad_norm": 0.27976104617118835,
120
+ "learning_rate": 5e-05,
121
+ "loss": 0.8119,
122
+ "step": 160
123
+ },
124
+ {
125
+ "epoch": 0.2071907373552712,
126
+ "grad_norm": 0.290866881608963,
127
+ "learning_rate": 5e-05,
128
+ "loss": 0.8096,
129
+ "step": 170
130
+ },
131
+ {
132
+ "epoch": 0.21937842778793418,
133
+ "grad_norm": 0.29430949687957764,
134
+ "learning_rate": 5e-05,
135
+ "loss": 0.819,
136
+ "step": 180
137
+ },
138
+ {
139
+ "epoch": 0.2315661182205972,
140
+ "grad_norm": 0.27714136242866516,
141
+ "learning_rate": 5e-05,
142
+ "loss": 0.8187,
143
+ "step": 190
144
+ },
145
+ {
146
+ "epoch": 0.2437538086532602,
147
+ "grad_norm": 0.2997029423713684,
148
+ "learning_rate": 5e-05,
149
+ "loss": 0.808,
150
+ "step": 200
151
+ },
152
+ {
153
+ "epoch": 0.25594149908592323,
154
+ "grad_norm": 0.2539573609828949,
155
+ "learning_rate": 5e-05,
156
+ "loss": 0.7943,
157
+ "step": 210
158
+ },
159
+ {
160
+ "epoch": 0.2681291895185862,
161
+ "grad_norm": 0.23294171690940857,
162
+ "learning_rate": 5e-05,
163
+ "loss": 0.8188,
164
+ "step": 220
165
+ },
166
+ {
167
+ "epoch": 0.28031687995124926,
168
+ "grad_norm": 0.25016292929649353,
169
+ "learning_rate": 5e-05,
170
+ "loss": 0.8035,
171
+ "step": 230
172
+ },
173
+ {
174
+ "epoch": 0.29250457038391225,
175
+ "grad_norm": 0.37125322222709656,
176
+ "learning_rate": 5e-05,
177
+ "loss": 0.785,
178
+ "step": 240
179
+ },
180
+ {
181
+ "epoch": 0.30469226081657524,
182
+ "grad_norm": 0.21520452201366425,
183
+ "learning_rate": 5e-05,
184
+ "loss": 0.8088,
185
+ "step": 250
186
+ },
187
+ {
188
+ "epoch": 0.3168799512492383,
189
+ "grad_norm": 0.20990602672100067,
190
+ "learning_rate": 5e-05,
191
+ "loss": 0.7813,
192
+ "step": 260
193
+ },
194
+ {
195
+ "epoch": 0.3290676416819013,
196
+ "grad_norm": 0.24193961918354034,
197
+ "learning_rate": 5e-05,
198
+ "loss": 0.7979,
199
+ "step": 270
200
+ },
201
+ {
202
+ "epoch": 0.34125533211456427,
203
+ "grad_norm": 0.22766050696372986,
204
+ "learning_rate": 5e-05,
205
+ "loss": 0.785,
206
+ "step": 280
207
+ },
208
+ {
209
+ "epoch": 0.3534430225472273,
210
+ "grad_norm": 0.2662702202796936,
211
+ "learning_rate": 5e-05,
212
+ "loss": 0.7766,
213
+ "step": 290
214
+ },
215
+ {
216
+ "epoch": 0.3656307129798903,
217
+ "grad_norm": 0.23020495474338531,
218
+ "learning_rate": 5e-05,
219
+ "loss": 0.7687,
220
+ "step": 300
221
+ },
222
+ {
223
+ "epoch": 0.3778184034125533,
224
+ "grad_norm": 0.21770520508289337,
225
+ "learning_rate": 5e-05,
226
+ "loss": 0.7815,
227
+ "step": 310
228
+ },
229
+ {
230
+ "epoch": 0.39000609384521634,
231
+ "grad_norm": 0.22677403688430786,
232
+ "learning_rate": 5e-05,
233
+ "loss": 0.7999,
234
+ "step": 320
235
+ },
236
+ {
237
+ "epoch": 0.40219378427787933,
238
+ "grad_norm": 0.24891860783100128,
239
+ "learning_rate": 5e-05,
240
+ "loss": 0.7946,
241
+ "step": 330
242
+ },
243
+ {
244
+ "epoch": 0.4143814747105424,
245
+ "grad_norm": 0.2506854832172394,
246
+ "learning_rate": 5e-05,
247
+ "loss": 0.7889,
248
+ "step": 340
249
+ },
250
+ {
251
+ "epoch": 0.42656916514320536,
252
+ "grad_norm": 0.24388180673122406,
253
+ "learning_rate": 5e-05,
254
+ "loss": 0.79,
255
+ "step": 350
256
+ },
257
+ {
258
+ "epoch": 0.43875685557586835,
259
+ "grad_norm": 0.22329650819301605,
260
+ "learning_rate": 5e-05,
261
+ "loss": 0.7632,
262
+ "step": 360
263
+ },
264
+ {
265
+ "epoch": 0.4509445460085314,
266
+ "grad_norm": 0.2240927815437317,
267
+ "learning_rate": 5e-05,
268
+ "loss": 0.8042,
269
+ "step": 370
270
+ },
271
+ {
272
+ "epoch": 0.4631322364411944,
273
+ "grad_norm": 0.2578203082084656,
274
+ "learning_rate": 5e-05,
275
+ "loss": 0.7667,
276
+ "step": 380
277
+ },
278
+ {
279
+ "epoch": 0.4753199268738574,
280
+ "grad_norm": 0.25349217653274536,
281
+ "learning_rate": 5e-05,
282
+ "loss": 0.7911,
283
+ "step": 390
284
+ },
285
+ {
286
+ "epoch": 0.4875076173065204,
287
+ "grad_norm": 0.254446804523468,
288
+ "learning_rate": 5e-05,
289
+ "loss": 0.7909,
290
+ "step": 400
291
+ },
292
+ {
293
+ "epoch": 0.4996953077391834,
294
+ "grad_norm": 0.24253171682357788,
295
+ "learning_rate": 5e-05,
296
+ "loss": 0.7896,
297
+ "step": 410
298
+ },
299
+ {
300
+ "epoch": 0.5118829981718465,
301
+ "grad_norm": 0.27901458740234375,
302
+ "learning_rate": 5e-05,
303
+ "loss": 0.7828,
304
+ "step": 420
305
+ },
306
+ {
307
+ "epoch": 0.5240706886045094,
308
+ "grad_norm": 0.31884434819221497,
309
+ "learning_rate": 5e-05,
310
+ "loss": 0.7919,
311
+ "step": 430
312
+ },
313
+ {
314
+ "epoch": 0.5362583790371724,
315
+ "grad_norm": 0.22920799255371094,
316
+ "learning_rate": 5e-05,
317
+ "loss": 0.7777,
318
+ "step": 440
319
+ },
320
+ {
321
+ "epoch": 0.5484460694698354,
322
+ "grad_norm": 0.22606618702411652,
323
+ "learning_rate": 5e-05,
324
+ "loss": 0.7661,
325
+ "step": 450
326
+ },
327
+ {
328
+ "epoch": 0.5606337599024985,
329
+ "grad_norm": 0.21604052186012268,
330
+ "learning_rate": 5e-05,
331
+ "loss": 0.769,
332
+ "step": 460
333
+ },
334
+ {
335
+ "epoch": 0.5728214503351615,
336
+ "grad_norm": 0.2296101152896881,
337
+ "learning_rate": 5e-05,
338
+ "loss": 0.7626,
339
+ "step": 470
340
+ },
341
+ {
342
+ "epoch": 0.5850091407678245,
343
+ "grad_norm": 0.24162150919437408,
344
+ "learning_rate": 5e-05,
345
+ "loss": 0.7831,
346
+ "step": 480
347
+ },
348
+ {
349
+ "epoch": 0.5971968312004875,
350
+ "grad_norm": 0.22364799678325653,
351
+ "learning_rate": 5e-05,
352
+ "loss": 0.7737,
353
+ "step": 490
354
+ },
355
+ {
356
+ "epoch": 0.6093845216331505,
357
+ "grad_norm": 0.23953403532505035,
358
+ "learning_rate": 5e-05,
359
+ "loss": 0.787,
360
+ "step": 500
361
+ },
362
+ {
363
+ "epoch": 0.6215722120658135,
364
+ "grad_norm": 0.22349311411380768,
365
+ "learning_rate": 5e-05,
366
+ "loss": 0.7763,
367
+ "step": 510
368
+ },
369
+ {
370
+ "epoch": 0.6337599024984766,
371
+ "grad_norm": 0.21996377408504486,
372
+ "learning_rate": 5e-05,
373
+ "loss": 0.7596,
374
+ "step": 520
375
+ },
376
+ {
377
+ "epoch": 0.6459475929311396,
378
+ "grad_norm": 0.245538130402565,
379
+ "learning_rate": 5e-05,
380
+ "loss": 0.7785,
381
+ "step": 530
382
+ },
383
+ {
384
+ "epoch": 0.6581352833638026,
385
+ "grad_norm": 0.23442520201206207,
386
+ "learning_rate": 5e-05,
387
+ "loss": 0.7633,
388
+ "step": 540
389
+ },
390
+ {
391
+ "epoch": 0.6703229737964655,
392
+ "grad_norm": 0.22274106740951538,
393
+ "learning_rate": 5e-05,
394
+ "loss": 0.7756,
395
+ "step": 550
396
+ },
397
+ {
398
+ "epoch": 0.6825106642291285,
399
+ "grad_norm": 0.21717530488967896,
400
+ "learning_rate": 5e-05,
401
+ "loss": 0.7726,
402
+ "step": 560
403
+ },
404
+ {
405
+ "epoch": 0.6946983546617916,
406
+ "grad_norm": 0.2465955913066864,
407
+ "learning_rate": 5e-05,
408
+ "loss": 0.7476,
409
+ "step": 570
410
+ },
411
+ {
412
+ "epoch": 0.7068860450944546,
413
+ "grad_norm": 0.2318575084209442,
414
+ "learning_rate": 5e-05,
415
+ "loss": 0.7584,
416
+ "step": 580
417
+ },
418
+ {
419
+ "epoch": 0.7190737355271176,
420
+ "grad_norm": 0.25114378333091736,
421
+ "learning_rate": 5e-05,
422
+ "loss": 0.7625,
423
+ "step": 590
424
+ },
425
+ {
426
+ "epoch": 0.7312614259597806,
427
+ "grad_norm": 0.19145755469799042,
428
+ "learning_rate": 5e-05,
429
+ "loss": 0.7577,
430
+ "step": 600
431
+ },
432
+ {
433
+ "epoch": 0.7434491163924436,
434
+ "grad_norm": 0.19129066169261932,
435
+ "learning_rate": 5e-05,
436
+ "loss": 0.7772,
437
+ "step": 610
438
+ },
439
+ {
440
+ "epoch": 0.7556368068251066,
441
+ "grad_norm": 0.2014475166797638,
442
+ "learning_rate": 5e-05,
443
+ "loss": 0.7601,
444
+ "step": 620
445
+ },
446
+ {
447
+ "epoch": 0.7678244972577697,
448
+ "grad_norm": 0.21480509638786316,
449
+ "learning_rate": 5e-05,
450
+ "loss": 0.7665,
451
+ "step": 630
452
+ },
453
+ {
454
+ "epoch": 0.7800121876904327,
455
+ "grad_norm": 0.2633228302001953,
456
+ "learning_rate": 5e-05,
457
+ "loss": 0.7565,
458
+ "step": 640
459
+ },
460
+ {
461
+ "epoch": 0.7921998781230957,
462
+ "grad_norm": 0.24633203446865082,
463
+ "learning_rate": 5e-05,
464
+ "loss": 0.7602,
465
+ "step": 650
466
+ },
467
+ {
468
+ "epoch": 0.8043875685557587,
469
+ "grad_norm": 0.19770534336566925,
470
+ "learning_rate": 5e-05,
471
+ "loss": 0.7685,
472
+ "step": 660
473
+ },
474
+ {
475
+ "epoch": 0.8165752589884216,
476
+ "grad_norm": 0.2267930805683136,
477
+ "learning_rate": 5e-05,
478
+ "loss": 0.7516,
479
+ "step": 670
480
+ },
481
+ {
482
+ "epoch": 0.8287629494210847,
483
+ "grad_norm": 0.238111674785614,
484
+ "learning_rate": 5e-05,
485
+ "loss": 0.7487,
486
+ "step": 680
487
+ },
488
+ {
489
+ "epoch": 0.8409506398537477,
490
+ "grad_norm": 0.19528904557228088,
491
+ "learning_rate": 5e-05,
492
+ "loss": 0.7467,
493
+ "step": 690
494
+ },
495
+ {
496
+ "epoch": 0.8531383302864107,
497
+ "grad_norm": 0.21305914223194122,
498
+ "learning_rate": 5e-05,
499
+ "loss": 0.7527,
500
+ "step": 700
501
+ },
502
+ {
503
+ "epoch": 0.8653260207190737,
504
+ "grad_norm": 0.2205205112695694,
505
+ "learning_rate": 5e-05,
506
+ "loss": 0.7639,
507
+ "step": 710
508
+ },
509
+ {
510
+ "epoch": 0.8775137111517367,
511
+ "grad_norm": 0.22429127991199493,
512
+ "learning_rate": 5e-05,
513
+ "loss": 0.7744,
514
+ "step": 720
515
+ },
516
+ {
517
+ "epoch": 0.8897014015843998,
518
+ "grad_norm": 0.23755772411823273,
519
+ "learning_rate": 5e-05,
520
+ "loss": 0.7446,
521
+ "step": 730
522
+ },
523
+ {
524
+ "epoch": 0.9018890920170628,
525
+ "grad_norm": 0.17592693865299225,
526
+ "learning_rate": 5e-05,
527
+ "loss": 0.7442,
528
+ "step": 740
529
+ },
530
+ {
531
+ "epoch": 0.9140767824497258,
532
+ "grad_norm": 0.2266898900270462,
533
+ "learning_rate": 5e-05,
534
+ "loss": 0.7326,
535
+ "step": 750
536
+ },
537
+ {
538
+ "epoch": 0.9262644728823888,
539
+ "grad_norm": 0.23541569709777832,
540
+ "learning_rate": 5e-05,
541
+ "loss": 0.7338,
542
+ "step": 760
543
+ },
544
+ {
545
+ "epoch": 0.9384521633150518,
546
+ "grad_norm": 0.22233888506889343,
547
+ "learning_rate": 5e-05,
548
+ "loss": 0.7385,
549
+ "step": 770
550
+ },
551
+ {
552
+ "epoch": 0.9506398537477148,
553
+ "grad_norm": 0.23714767396450043,
554
+ "learning_rate": 5e-05,
555
+ "loss": 0.7306,
556
+ "step": 780
557
+ },
558
+ {
559
+ "epoch": 0.9628275441803779,
560
+ "grad_norm": 0.21191060543060303,
561
+ "learning_rate": 5e-05,
562
+ "loss": 0.736,
563
+ "step": 790
564
+ },
565
+ {
566
+ "epoch": 0.9750152346130408,
567
+ "grad_norm": 0.22759106755256653,
568
+ "learning_rate": 5e-05,
569
+ "loss": 0.7505,
570
+ "step": 800
571
+ },
572
+ {
573
+ "epoch": 0.9872029250457038,
574
+ "grad_norm": 0.19895903766155243,
575
+ "learning_rate": 5e-05,
576
+ "loss": 0.7536,
577
+ "step": 810
578
+ },
579
+ {
580
+ "epoch": 0.9993906154783668,
581
+ "grad_norm": 0.3061685860157013,
582
+ "learning_rate": 5e-05,
583
+ "loss": 0.7493,
584
+ "step": 820
585
+ },
586
+ {
587
+ "epoch": 1.0109689213893966,
588
+ "grad_norm": 0.18868015706539154,
589
+ "learning_rate": 5e-05,
590
+ "loss": 0.7342,
591
+ "step": 830
592
+ },
593
+ {
594
+ "epoch": 1.0231566118220596,
595
+ "grad_norm": 0.1797628402709961,
596
+ "learning_rate": 5e-05,
597
+ "loss": 0.7192,
598
+ "step": 840
599
+ },
600
+ {
601
+ "epoch": 1.0353443022547226,
602
+ "grad_norm": 0.19511596858501434,
603
+ "learning_rate": 5e-05,
604
+ "loss": 0.7002,
605
+ "step": 850
606
+ },
607
+ {
608
+ "epoch": 1.0475319926873858,
609
+ "grad_norm": 0.2123185396194458,
610
+ "learning_rate": 5e-05,
611
+ "loss": 0.7236,
612
+ "step": 860
613
+ },
614
+ {
615
+ "epoch": 1.0597196831200488,
616
+ "grad_norm": 0.2047470211982727,
617
+ "learning_rate": 5e-05,
618
+ "loss": 0.711,
619
+ "step": 870
620
+ },
621
+ {
622
+ "epoch": 1.0719073735527118,
623
+ "grad_norm": 0.18809868395328522,
624
+ "learning_rate": 5e-05,
625
+ "loss": 0.7253,
626
+ "step": 880
627
+ },
628
+ {
629
+ "epoch": 1.0840950639853748,
630
+ "grad_norm": 0.2207012176513672,
631
+ "learning_rate": 5e-05,
632
+ "loss": 0.6971,
633
+ "step": 890
634
+ },
635
+ {
636
+ "epoch": 1.0962827544180378,
637
+ "grad_norm": 0.1912035495042801,
638
+ "learning_rate": 5e-05,
639
+ "loss": 0.7239,
640
+ "step": 900
641
+ },
642
+ {
643
+ "epoch": 1.1084704448507008,
644
+ "grad_norm": 0.23486360907554626,
645
+ "learning_rate": 5e-05,
646
+ "loss": 0.7335,
647
+ "step": 910
648
+ },
649
+ {
650
+ "epoch": 1.1206581352833638,
651
+ "grad_norm": 0.22189787030220032,
652
+ "learning_rate": 5e-05,
653
+ "loss": 0.7149,
654
+ "step": 920
655
+ },
656
+ {
657
+ "epoch": 1.1328458257160268,
658
+ "grad_norm": 0.25818824768066406,
659
+ "learning_rate": 5e-05,
660
+ "loss": 0.7322,
661
+ "step": 930
662
+ },
663
+ {
664
+ "epoch": 1.1450335161486898,
665
+ "grad_norm": 0.2088775783777237,
666
+ "learning_rate": 5e-05,
667
+ "loss": 0.7382,
668
+ "step": 940
669
+ },
670
+ {
671
+ "epoch": 1.1572212065813527,
672
+ "grad_norm": 0.19445833563804626,
673
+ "learning_rate": 5e-05,
674
+ "loss": 0.7223,
675
+ "step": 950
676
+ },
677
+ {
678
+ "epoch": 1.1694088970140157,
679
+ "grad_norm": 0.21489456295967102,
680
+ "learning_rate": 5e-05,
681
+ "loss": 0.7317,
682
+ "step": 960
683
+ },
684
+ {
685
+ "epoch": 1.181596587446679,
686
+ "grad_norm": 0.24050794541835785,
687
+ "learning_rate": 5e-05,
688
+ "loss": 0.7275,
689
+ "step": 970
690
+ },
691
+ {
692
+ "epoch": 1.193784277879342,
693
+ "grad_norm": 0.18142102658748627,
694
+ "learning_rate": 5e-05,
695
+ "loss": 0.7046,
696
+ "step": 980
697
+ },
698
+ {
699
+ "epoch": 1.205971968312005,
700
+ "grad_norm": 0.4781077802181244,
701
+ "learning_rate": 5e-05,
702
+ "loss": 0.7433,
703
+ "step": 990
704
+ },
705
+ {
706
+ "epoch": 1.218159658744668,
707
+ "grad_norm": 0.20628023147583008,
708
+ "learning_rate": 5e-05,
709
+ "loss": 0.7403,
710
+ "step": 1000
711
+ },
712
+ {
713
+ "epoch": 1.230347349177331,
714
+ "grad_norm": 0.19117140769958496,
715
+ "learning_rate": 5e-05,
716
+ "loss": 0.7308,
717
+ "step": 1010
718
+ },
719
+ {
720
+ "epoch": 1.242535039609994,
721
+ "grad_norm": 0.2107745110988617,
722
+ "learning_rate": 5e-05,
723
+ "loss": 0.7201,
724
+ "step": 1020
725
+ },
726
+ {
727
+ "epoch": 1.2547227300426569,
728
+ "grad_norm": 0.20551154017448425,
729
+ "learning_rate": 5e-05,
730
+ "loss": 0.7038,
731
+ "step": 1030
732
+ },
733
+ {
734
+ "epoch": 1.2669104204753199,
735
+ "grad_norm": 0.26577022671699524,
736
+ "learning_rate": 5e-05,
737
+ "loss": 0.7096,
738
+ "step": 1040
739
+ },
740
+ {
741
+ "epoch": 1.2790981109079829,
742
+ "grad_norm": 0.22094209492206573,
743
+ "learning_rate": 5e-05,
744
+ "loss": 0.718,
745
+ "step": 1050
746
+ },
747
+ {
748
+ "epoch": 1.291285801340646,
749
+ "grad_norm": 0.18224076926708221,
750
+ "learning_rate": 5e-05,
751
+ "loss": 0.7261,
752
+ "step": 1060
753
+ },
754
+ {
755
+ "epoch": 1.303473491773309,
756
+ "grad_norm": 0.226251482963562,
757
+ "learning_rate": 5e-05,
758
+ "loss": 0.7372,
759
+ "step": 1070
760
+ },
761
+ {
762
+ "epoch": 1.315661182205972,
763
+ "grad_norm": 0.20186461508274078,
764
+ "learning_rate": 5e-05,
765
+ "loss": 0.7074,
766
+ "step": 1080
767
+ },
768
+ {
769
+ "epoch": 1.327848872638635,
770
+ "grad_norm": 0.22229762375354767,
771
+ "learning_rate": 5e-05,
772
+ "loss": 0.7069,
773
+ "step": 1090
774
+ },
775
+ {
776
+ "epoch": 1.340036563071298,
777
+ "grad_norm": 0.1776452213525772,
778
+ "learning_rate": 5e-05,
779
+ "loss": 0.714,
780
+ "step": 1100
781
+ },
782
+ {
783
+ "epoch": 1.352224253503961,
784
+ "grad_norm": 0.20680229365825653,
785
+ "learning_rate": 5e-05,
786
+ "loss": 0.7066,
787
+ "step": 1110
788
+ },
789
+ {
790
+ "epoch": 1.364411943936624,
791
+ "grad_norm": 0.21842031180858612,
792
+ "learning_rate": 5e-05,
793
+ "loss": 0.7281,
794
+ "step": 1120
795
+ },
796
+ {
797
+ "epoch": 1.376599634369287,
798
+ "grad_norm": 0.28037774562835693,
799
+ "learning_rate": 5e-05,
800
+ "loss": 0.7149,
801
+ "step": 1130
802
+ },
803
+ {
804
+ "epoch": 1.38878732480195,
805
+ "grad_norm": 0.20962920784950256,
806
+ "learning_rate": 5e-05,
807
+ "loss": 0.6951,
808
+ "step": 1140
809
+ },
810
+ {
811
+ "epoch": 1.400975015234613,
812
+ "grad_norm": 0.18751871585845947,
813
+ "learning_rate": 5e-05,
814
+ "loss": 0.7101,
815
+ "step": 1150
816
+ },
817
+ {
818
+ "epoch": 1.413162705667276,
819
+ "grad_norm": 0.2722964584827423,
820
+ "learning_rate": 5e-05,
821
+ "loss": 0.7205,
822
+ "step": 1160
823
+ },
824
+ {
825
+ "epoch": 1.425350396099939,
826
+ "grad_norm": 0.20932041108608246,
827
+ "learning_rate": 5e-05,
828
+ "loss": 0.7082,
829
+ "step": 1170
830
+ },
831
+ {
832
+ "epoch": 1.437538086532602,
833
+ "grad_norm": 0.17836257815361023,
834
+ "learning_rate": 5e-05,
835
+ "loss": 0.7232,
836
+ "step": 1180
837
+ },
838
+ {
839
+ "epoch": 1.449725776965265,
840
+ "grad_norm": 0.3180703818798065,
841
+ "learning_rate": 5e-05,
842
+ "loss": 0.7245,
843
+ "step": 1190
844
+ },
845
+ {
846
+ "epoch": 1.4619134673979282,
847
+ "grad_norm": 0.201424703001976,
848
+ "learning_rate": 5e-05,
849
+ "loss": 0.7232,
850
+ "step": 1200
851
+ },
852
+ {
853
+ "epoch": 1.4741011578305911,
854
+ "grad_norm": 0.20228999853134155,
855
+ "learning_rate": 5e-05,
856
+ "loss": 0.7169,
857
+ "step": 1210
858
+ },
859
+ {
860
+ "epoch": 1.4862888482632541,
861
+ "grad_norm": 0.17478394508361816,
862
+ "learning_rate": 5e-05,
863
+ "loss": 0.704,
864
+ "step": 1220
865
+ },
866
+ {
867
+ "epoch": 1.4984765386959171,
868
+ "grad_norm": 0.20061351358890533,
869
+ "learning_rate": 5e-05,
870
+ "loss": 0.6792,
871
+ "step": 1230
872
+ },
873
+ {
874
+ "epoch": 1.5106642291285801,
875
+ "grad_norm": 0.1857377588748932,
876
+ "learning_rate": 5e-05,
877
+ "loss": 0.709,
878
+ "step": 1240
879
+ },
880
+ {
881
+ "epoch": 1.522851919561243,
882
+ "grad_norm": 0.19873379170894623,
883
+ "learning_rate": 5e-05,
884
+ "loss": 0.7007,
885
+ "step": 1250
886
+ },
887
+ {
888
+ "epoch": 1.5350396099939063,
889
+ "grad_norm": 0.19807074964046478,
890
+ "learning_rate": 5e-05,
891
+ "loss": 0.7122,
892
+ "step": 1260
893
+ },
894
+ {
895
+ "epoch": 1.5472273004265693,
896
+ "grad_norm": 0.16746796667575836,
897
+ "learning_rate": 5e-05,
898
+ "loss": 0.7005,
899
+ "step": 1270
900
+ },
901
+ {
902
+ "epoch": 1.5594149908592323,
903
+ "grad_norm": 0.2750326991081238,
904
+ "learning_rate": 5e-05,
905
+ "loss": 0.7166,
906
+ "step": 1280
907
+ },
908
+ {
909
+ "epoch": 1.5716026812918953,
910
+ "grad_norm": 0.2776738107204437,
911
+ "learning_rate": 5e-05,
912
+ "loss": 0.7405,
913
+ "step": 1290
914
+ },
915
+ {
916
+ "epoch": 1.5837903717245583,
917
+ "grad_norm": 0.3494517207145691,
918
+ "learning_rate": 5e-05,
919
+ "loss": 0.7192,
920
+ "step": 1300
921
+ },
922
+ {
923
+ "epoch": 1.5959780621572213,
924
+ "grad_norm": 0.25028494000434875,
925
+ "learning_rate": 5e-05,
926
+ "loss": 0.741,
927
+ "step": 1310
928
+ },
929
+ {
930
+ "epoch": 1.6081657525898843,
931
+ "grad_norm": 0.23914143443107605,
932
+ "learning_rate": 5e-05,
933
+ "loss": 0.72,
934
+ "step": 1320
935
+ },
936
+ {
937
+ "epoch": 1.6203534430225472,
938
+ "grad_norm": 0.2052862048149109,
939
+ "learning_rate": 5e-05,
940
+ "loss": 0.7112,
941
+ "step": 1330
942
+ },
943
+ {
944
+ "epoch": 1.6325411334552102,
945
+ "grad_norm": 0.22285936772823334,
946
+ "learning_rate": 5e-05,
947
+ "loss": 0.7165,
948
+ "step": 1340
949
+ },
950
+ {
951
+ "epoch": 1.6447288238878732,
952
+ "grad_norm": 0.20810848474502563,
953
+ "learning_rate": 5e-05,
954
+ "loss": 0.7122,
955
+ "step": 1350
956
+ },
957
+ {
958
+ "epoch": 1.6569165143205362,
959
+ "grad_norm": 0.20422884821891785,
960
+ "learning_rate": 5e-05,
961
+ "loss": 0.7082,
962
+ "step": 1360
963
+ },
964
+ {
965
+ "epoch": 1.6691042047531992,
966
+ "grad_norm": 0.18164314329624176,
967
+ "learning_rate": 5e-05,
968
+ "loss": 0.7131,
969
+ "step": 1370
970
+ },
971
+ {
972
+ "epoch": 1.6812918951858622,
973
+ "grad_norm": 0.19893014430999756,
974
+ "learning_rate": 5e-05,
975
+ "loss": 0.7055,
976
+ "step": 1380
977
+ },
978
+ {
979
+ "epoch": 1.6934795856185252,
980
+ "grad_norm": 0.24865971505641937,
981
+ "learning_rate": 5e-05,
982
+ "loss": 0.6994,
983
+ "step": 1390
984
+ },
985
+ {
986
+ "epoch": 1.7056672760511882,
987
+ "grad_norm": 0.18078123033046722,
988
+ "learning_rate": 5e-05,
989
+ "loss": 0.7162,
990
+ "step": 1400
991
+ },
992
+ {
993
+ "epoch": 1.7178549664838512,
994
+ "grad_norm": 0.1973421722650528,
995
+ "learning_rate": 5e-05,
996
+ "loss": 0.7136,
997
+ "step": 1410
998
+ },
999
+ {
1000
+ "epoch": 1.7300426569165142,
1001
+ "grad_norm": 0.2057911604642868,
1002
+ "learning_rate": 5e-05,
1003
+ "loss": 0.7004,
1004
+ "step": 1420
1005
+ },
1006
+ {
1007
+ "epoch": 1.7422303473491774,
1008
+ "grad_norm": 0.18640495836734772,
1009
+ "learning_rate": 5e-05,
1010
+ "loss": 0.7102,
1011
+ "step": 1430
1012
+ }
1013
+ ],
1014
+ "logging_steps": 10,
1015
+ "max_steps": 2463,
1016
+ "num_input_tokens_seen": 0,
1017
+ "num_train_epochs": 3,
1018
+ "save_steps": 205,
1019
+ "stateful_callbacks": {
1020
+ "TrainerControl": {
1021
+ "args": {
1022
+ "should_epoch_stop": false,
1023
+ "should_evaluate": false,
1024
+ "should_log": false,
1025
+ "should_save": true,
1026
+ "should_training_stop": false
1027
+ },
1028
+ "attributes": {}
1029
+ }
1030
+ },
1031
+ "total_flos": 5.8872689779519e+18,
1032
+ "train_batch_size": 4,
1033
+ "trial_name": null,
1034
+ "trial_params": null
1035
+ }
openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1435/training_args.bin ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:56896ea21940e33d7e2934267dff2a5b4a4a80195a0f165245f6617a0f88b90b
3
+ size 7377
openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1435/zero_to_fp32.py ADDED
@@ -0,0 +1,760 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python
2
+
3
+ # Copyright (c) Microsoft Corporation.
4
+ # SPDX-License-Identifier: Apache-2.0
5
+
6
+ # DeepSpeed Team
7
+
8
+ # This script extracts fp32 consolidated weights from a zero 1, 2 and 3 DeepSpeed checkpoints. It gets
9
+ # copied into the top level checkpoint dir, so the user can easily do the conversion at any point in
10
+ # the future. Once extracted, the weights don't require DeepSpeed and can be used in any
11
+ # application.
12
+ #
13
+ # example:
14
+ # python zero_to_fp32.py . output_dir/
15
+ # or
16
+ # python zero_to_fp32.py . output_dir/ --safe_serialization
17
+
18
+ import argparse
19
+ import torch
20
+ import glob
21
+ import math
22
+ import os
23
+ import re
24
+ import gc
25
+ import json
26
+ import numpy as np
27
+ from tqdm import tqdm
28
+ from collections import OrderedDict
29
+ from dataclasses import dataclass
30
+
31
+ # while this script doesn't use deepspeed to recover data, since the checkpoints are pickled with
32
+ # DeepSpeed data structures it has to be available in the current python environment.
33
+ from deepspeed.utils import logger
34
+ from deepspeed.checkpoint.constants import (DS_VERSION, OPTIMIZER_STATE_DICT, SINGLE_PARTITION_OF_FP32_GROUPS,
35
+ FP32_FLAT_GROUPS, ZERO_STAGE, PARTITION_COUNT, PARAM_SHAPES, BUFFER_NAMES,
36
+ FROZEN_PARAM_SHAPES, FROZEN_PARAM_FRAGMENTS)
37
+
38
+
39
+ @dataclass
40
+ class zero_model_state:
41
+ buffers: dict()
42
+ param_shapes: dict()
43
+ shared_params: list
44
+ ds_version: int
45
+ frozen_param_shapes: dict()
46
+ frozen_param_fragments: dict()
47
+
48
+
49
+ debug = 0
50
+
51
+ # load to cpu
52
+ device = torch.device('cpu')
53
+
54
+
55
+ def atoi(text):
56
+ return int(text) if text.isdigit() else text
57
+
58
+
59
+ def natural_keys(text):
60
+ '''
61
+ alist.sort(key=natural_keys) sorts in human order
62
+ http://nedbatchelder.com/blog/200712/human_sorting.html
63
+ (See Toothy's implementation in the comments)
64
+ '''
65
+ return [atoi(c) for c in re.split(r'(\d+)', text)]
66
+
67
+
68
+ def get_model_state_file(checkpoint_dir, zero_stage):
69
+ if not os.path.isdir(checkpoint_dir):
70
+ raise FileNotFoundError(f"Directory '{checkpoint_dir}' doesn't exist")
71
+
72
+ # there should be only one file
73
+ if zero_stage <= 2:
74
+ file = os.path.join(checkpoint_dir, "mp_rank_00_model_states.pt")
75
+ elif zero_stage == 3:
76
+ file = os.path.join(checkpoint_dir, "zero_pp_rank_0_mp_rank_00_model_states.pt")
77
+
78
+ if not os.path.exists(file):
79
+ raise FileNotFoundError(f"can't find model states file at '{file}'")
80
+
81
+ return file
82
+
83
+
84
+ def get_checkpoint_files(checkpoint_dir, glob_pattern):
85
+ # XXX: need to test that this simple glob rule works for multi-node setup too
86
+ ckpt_files = sorted(glob.glob(os.path.join(checkpoint_dir, glob_pattern)), key=natural_keys)
87
+
88
+ if len(ckpt_files) == 0:
89
+ raise FileNotFoundError(f"can't find {glob_pattern} files in directory '{checkpoint_dir}'")
90
+
91
+ return ckpt_files
92
+
93
+
94
+ def get_optim_files(checkpoint_dir):
95
+ return get_checkpoint_files(checkpoint_dir, "*_optim_states.pt")
96
+
97
+
98
+ def get_model_state_files(checkpoint_dir):
99
+ return get_checkpoint_files(checkpoint_dir, "*_model_states.pt")
100
+
101
+
102
+ def parse_model_states(files):
103
+ zero_model_states = []
104
+ for file in files:
105
+ state_dict = torch.load(file, map_location=device, weights_only=False)
106
+
107
+ if BUFFER_NAMES not in state_dict:
108
+ raise ValueError(f"{file} is not a model state checkpoint")
109
+ buffer_names = state_dict[BUFFER_NAMES]
110
+ if debug:
111
+ print("Found buffers:", buffer_names)
112
+
113
+ # recover just the buffers while restoring them to fp32 if they were saved in fp16
114
+ buffers = {k: v.float() for k, v in state_dict["module"].items() if k in buffer_names}
115
+ param_shapes = state_dict[PARAM_SHAPES]
116
+
117
+ # collect parameters that are included in param_shapes
118
+ param_names = []
119
+ for s in param_shapes:
120
+ for name in s.keys():
121
+ param_names.append(name)
122
+
123
+ # update with frozen parameters
124
+ frozen_param_shapes = state_dict.get(FROZEN_PARAM_SHAPES, None)
125
+ if frozen_param_shapes is not None:
126
+ if debug:
127
+ print(f"Found frozen_param_shapes: {frozen_param_shapes}")
128
+ param_names += list(frozen_param_shapes.keys())
129
+
130
+ # handle shared params
131
+ shared_params = [[k, v] for k, v in state_dict["shared_params"].items()]
132
+
133
+ ds_version = state_dict.get(DS_VERSION, None)
134
+
135
+ frozen_param_fragments = state_dict.get(FROZEN_PARAM_FRAGMENTS, None)
136
+
137
+ z_model_state = zero_model_state(buffers=buffers,
138
+ param_shapes=param_shapes,
139
+ shared_params=shared_params,
140
+ ds_version=ds_version,
141
+ frozen_param_shapes=frozen_param_shapes,
142
+ frozen_param_fragments=frozen_param_fragments)
143
+ zero_model_states.append(z_model_state)
144
+
145
+ return zero_model_states
146
+
147
+
148
+ def parse_optim_states(files, ds_checkpoint_dir):
149
+ total_files = len(files)
150
+ state_dicts = []
151
+ for f in tqdm(files, desc='Loading checkpoint shards'):
152
+ state_dict = torch.load(f, map_location=device, mmap=True, weights_only=False)
153
+ # immediately discard the potentially huge 2 optimizer states as we only care for fp32 master weights
154
+ # and also handle the case where it was already removed by another helper script
155
+ state_dict["optimizer_state_dict"].pop("optimizer_state_dict", None)
156
+ state_dicts.append(state_dict)
157
+
158
+ if ZERO_STAGE not in state_dicts[0][OPTIMIZER_STATE_DICT]:
159
+ raise ValueError(f"{files[0]} is not a zero checkpoint")
160
+ zero_stage = state_dicts[0][OPTIMIZER_STATE_DICT][ZERO_STAGE]
161
+ world_size = state_dicts[0][OPTIMIZER_STATE_DICT][PARTITION_COUNT]
162
+
163
+ # For ZeRO-2 each param group can have different partition_count as data parallelism for expert
164
+ # parameters can be different from data parallelism for non-expert parameters. So we can just
165
+ # use the max of the partition_count to get the dp world_size.
166
+
167
+ if type(world_size) is list:
168
+ world_size = max(world_size)
169
+
170
+ if world_size != total_files:
171
+ raise ValueError(
172
+ f"Expected {world_size} of '*_optim_states.pt' under '{ds_checkpoint_dir}' but found {total_files} files. "
173
+ "Possibly due to an overwrite of an old checkpoint, or a checkpoint didn't get saved by one or more processes."
174
+ )
175
+
176
+ # the groups are named differently in each stage
177
+ if zero_stage <= 2:
178
+ fp32_groups_key = SINGLE_PARTITION_OF_FP32_GROUPS
179
+ elif zero_stage == 3:
180
+ fp32_groups_key = FP32_FLAT_GROUPS
181
+ else:
182
+ raise ValueError(f"unknown zero stage {zero_stage}")
183
+
184
+ fp32_flat_groups = [state_dicts[i][OPTIMIZER_STATE_DICT][fp32_groups_key] for i in range(len(state_dicts))]
185
+ return zero_stage, world_size, fp32_flat_groups
186
+
187
+
188
+ def _get_fp32_state_dict_from_zero_checkpoint(ds_checkpoint_dir, exclude_frozen_parameters):
189
+ """
190
+ Returns fp32 state_dict reconstructed from ds checkpoint
191
+
192
+ Args:
193
+ - ``ds_checkpoint_dir``: path to the deepspeed checkpoint folder (where the optimizer files are)
194
+
195
+ """
196
+ print(f"Processing zero checkpoint '{ds_checkpoint_dir}'")
197
+
198
+ optim_files = get_optim_files(ds_checkpoint_dir)
199
+ zero_stage, world_size, fp32_flat_groups = parse_optim_states(optim_files, ds_checkpoint_dir)
200
+ print(f"Detected checkpoint of type zero stage {zero_stage}, world_size: {world_size}")
201
+
202
+ model_files = get_model_state_files(ds_checkpoint_dir)
203
+
204
+ zero_model_states = parse_model_states(model_files)
205
+ print(f'Parsing checkpoint created by deepspeed=={zero_model_states[0].ds_version}')
206
+
207
+ if zero_stage <= 2:
208
+ return _get_fp32_state_dict_from_zero2_checkpoint(world_size, fp32_flat_groups, zero_model_states,
209
+ exclude_frozen_parameters)
210
+ elif zero_stage == 3:
211
+ return _get_fp32_state_dict_from_zero3_checkpoint(world_size, fp32_flat_groups, zero_model_states,
212
+ exclude_frozen_parameters)
213
+
214
+
215
+ def _zero2_merge_frozen_params(state_dict, zero_model_states):
216
+ if zero_model_states[0].frozen_param_shapes is None or len(zero_model_states[0].frozen_param_shapes) == 0:
217
+ return
218
+
219
+ frozen_param_shapes = zero_model_states[0].frozen_param_shapes
220
+ frozen_param_fragments = zero_model_states[0].frozen_param_fragments
221
+
222
+ if debug:
223
+ num_elem = sum(s.numel() for s in frozen_param_shapes.values())
224
+ print(f'rank 0: {FROZEN_PARAM_SHAPES}.numel = {num_elem}')
225
+
226
+ wanted_params = len(frozen_param_shapes)
227
+ wanted_numel = sum(s.numel() for s in frozen_param_shapes.values())
228
+ avail_numel = sum([p.numel() for p in frozen_param_fragments.values()])
229
+ print(f'Frozen params: Have {avail_numel} numels to process.')
230
+ print(f'Frozen params: Need {wanted_numel} numels in {wanted_params} params')
231
+
232
+ total_params = 0
233
+ total_numel = 0
234
+ for name, shape in frozen_param_shapes.items():
235
+ total_params += 1
236
+ unpartitioned_numel = shape.numel()
237
+ total_numel += unpartitioned_numel
238
+
239
+ state_dict[name] = frozen_param_fragments[name]
240
+
241
+ if debug:
242
+ print(f"{name} full shape: {shape} unpartitioned numel {unpartitioned_numel} ")
243
+
244
+ print(f"Reconstructed Frozen fp32 state dict with {total_params} params {total_numel} elements")
245
+
246
+
247
+ def _has_callable(obj, fn):
248
+ attr = getattr(obj, fn, None)
249
+ return callable(attr)
250
+
251
+
252
+ def _zero2_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states):
253
+ param_shapes = zero_model_states[0].param_shapes
254
+
255
+ # Reconstruction protocol:
256
+ #
257
+ # XXX: document this
258
+
259
+ if debug:
260
+ for i in range(world_size):
261
+ for j in range(len(fp32_flat_groups[0])):
262
+ print(f"{FP32_FLAT_GROUPS}[{i}][{j}].shape={fp32_flat_groups[i][j].shape}")
263
+
264
+ # XXX: memory usage doubles here (zero2)
265
+ num_param_groups = len(fp32_flat_groups[0])
266
+ merged_single_partition_of_fp32_groups = []
267
+ for i in range(num_param_groups):
268
+ merged_partitions = [sd[i] for sd in fp32_flat_groups]
269
+ full_single_fp32_vector = torch.cat(merged_partitions, 0)
270
+ merged_single_partition_of_fp32_groups.append(full_single_fp32_vector)
271
+ avail_numel = sum(
272
+ [full_single_fp32_vector.numel() for full_single_fp32_vector in merged_single_partition_of_fp32_groups])
273
+
274
+ if debug:
275
+ wanted_params = sum([len(shapes) for shapes in param_shapes])
276
+ wanted_numel = sum([sum(shape.numel() for shape in shapes.values()) for shapes in param_shapes])
277
+ # not asserting if there is a mismatch due to possible padding
278
+ print(f"Have {avail_numel} numels to process.")
279
+ print(f"Need {wanted_numel} numels in {wanted_params} params.")
280
+
281
+ # params
282
+ # XXX: for huge models that can't fit into the host's RAM we will have to recode this to support
283
+ # out-of-core computing solution
284
+ total_numel = 0
285
+ total_params = 0
286
+ for shapes, full_single_fp32_vector in zip(param_shapes, merged_single_partition_of_fp32_groups):
287
+ offset = 0
288
+ avail_numel = full_single_fp32_vector.numel()
289
+ for name, shape in shapes.items():
290
+
291
+ unpartitioned_numel = shape.numel() if _has_callable(shape, 'numel') else math.prod(shape)
292
+ total_numel += unpartitioned_numel
293
+ total_params += 1
294
+
295
+ if debug:
296
+ print(f"{name} full shape: {shape} unpartitioned numel {unpartitioned_numel} ")
297
+ state_dict[name] = full_single_fp32_vector.narrow(0, offset, unpartitioned_numel).view(shape)
298
+ offset += unpartitioned_numel
299
+
300
+ # Z2 started to align to 2*world_size to improve nccl performance. Therefore both offset and
301
+ # avail_numel can differ by anywhere between 0..2*world_size. Due to two unrelated complex
302
+ # paddings performed in the code it's almost impossible to predict the exact numbers w/o the
303
+ # live optimizer object, so we are checking that the numbers are within the right range
304
+ align_to = 2 * world_size
305
+
306
+ def zero2_align(x):
307
+ return align_to * math.ceil(x / align_to)
308
+
309
+ if debug:
310
+ print(f"original offset={offset}, avail_numel={avail_numel}")
311
+
312
+ offset = zero2_align(offset)
313
+ avail_numel = zero2_align(avail_numel)
314
+
315
+ if debug:
316
+ print(f"aligned offset={offset}, avail_numel={avail_numel}")
317
+
318
+ # Sanity check
319
+ if offset != avail_numel:
320
+ raise ValueError(f"consumed {offset} numels out of {avail_numel} - something is wrong")
321
+
322
+ print(f"Reconstructed fp32 state dict with {total_params} params {total_numel} elements")
323
+
324
+
325
+ def _get_fp32_state_dict_from_zero2_checkpoint(world_size, fp32_flat_groups, zero_model_states,
326
+ exclude_frozen_parameters):
327
+ state_dict = OrderedDict()
328
+
329
+ # buffers
330
+ buffers = zero_model_states[0].buffers
331
+ state_dict.update(buffers)
332
+ if debug:
333
+ print(f"added {len(buffers)} buffers")
334
+
335
+ if not exclude_frozen_parameters:
336
+ _zero2_merge_frozen_params(state_dict, zero_model_states)
337
+
338
+ _zero2_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states)
339
+
340
+ # recover shared parameters
341
+ for pair in zero_model_states[0].shared_params:
342
+ if pair[1] in state_dict:
343
+ state_dict[pair[0]] = state_dict[pair[1]]
344
+
345
+ return state_dict
346
+
347
+
348
+ def zero3_partitioned_param_info(unpartitioned_numel, world_size):
349
+ remainder = unpartitioned_numel % world_size
350
+ padding_numel = (world_size - remainder) if remainder else 0
351
+ partitioned_numel = math.ceil(unpartitioned_numel / world_size)
352
+ return partitioned_numel, padding_numel
353
+
354
+
355
+ def _zero3_merge_frozen_params(state_dict, world_size, zero_model_states):
356
+ if zero_model_states[0].frozen_param_shapes is None or len(zero_model_states[0].frozen_param_shapes) == 0:
357
+ return
358
+
359
+ if debug:
360
+ for i in range(world_size):
361
+ num_elem = sum(s.numel() for s in zero_model_states[i].frozen_param_fragments.values())
362
+ print(f'rank {i}: {FROZEN_PARAM_SHAPES}.numel = {num_elem}')
363
+
364
+ frozen_param_shapes = zero_model_states[0].frozen_param_shapes
365
+ wanted_params = len(frozen_param_shapes)
366
+ wanted_numel = sum(s.numel() for s in frozen_param_shapes.values())
367
+ avail_numel = sum([p.numel() for p in zero_model_states[0].frozen_param_fragments.values()]) * world_size
368
+ print(f'Frozen params: Have {avail_numel} numels to process.')
369
+ print(f'Frozen params: Need {wanted_numel} numels in {wanted_params} params')
370
+
371
+ total_params = 0
372
+ total_numel = 0
373
+ for name, shape in zero_model_states[0].frozen_param_shapes.items():
374
+ total_params += 1
375
+ unpartitioned_numel = shape.numel()
376
+ total_numel += unpartitioned_numel
377
+
378
+ param_frags = tuple(model_state.frozen_param_fragments[name] for model_state in zero_model_states)
379
+ state_dict[name] = torch.cat(param_frags, 0).narrow(0, 0, unpartitioned_numel).view(shape)
380
+
381
+ partitioned_numel, partitioned_padding_numel = zero3_partitioned_param_info(unpartitioned_numel, world_size)
382
+
383
+ if debug:
384
+ print(
385
+ f"Frozen params: {total_params} {name} full shape: {shape} partition0 numel={partitioned_numel} partitioned_padding_numel={partitioned_padding_numel}"
386
+ )
387
+
388
+ print(f"Reconstructed Frozen fp32 state dict with {total_params} params {total_numel} elements")
389
+
390
+
391
+ class GatheredTensor:
392
+ """
393
+ A pseudo tensor that collects partitioned weights.
394
+ It is more memory efficient when there are multiple groups.
395
+ """
396
+
397
+ def __init__(self, flat_groups, flat_groups_offset, offset, partitioned_numel, shape):
398
+ self.flat_groups = flat_groups
399
+ self.flat_groups_offset = flat_groups_offset
400
+ self.offset = offset
401
+ self.partitioned_numel = partitioned_numel
402
+ self.shape = shape
403
+ self.dtype = self.flat_groups[0][0].dtype
404
+
405
+ def contiguous(self):
406
+ """
407
+ Merge partitioned weights from flat_groups into a single tensor.
408
+ """
409
+ end_idx = self.offset + self.partitioned_numel
410
+ world_size = len(self.flat_groups)
411
+ pad_flat_param_chunks = []
412
+
413
+ for rank_i in range(world_size):
414
+ # for each rank, we need to collect weights from related group/groups
415
+ flat_groups_at_rank_i = self.flat_groups[rank_i]
416
+ start_group_id = None
417
+ end_group_id = None
418
+ for group_id in range(len(self.flat_groups_offset)):
419
+ if self.flat_groups_offset[group_id] <= self.offset < self.flat_groups_offset[group_id + 1]:
420
+ start_group_id = group_id
421
+ if self.flat_groups_offset[group_id] < end_idx <= self.flat_groups_offset[group_id + 1]:
422
+ end_group_id = group_id
423
+ break
424
+ # collect weights from related group/groups
425
+ for group_id in range(start_group_id, end_group_id + 1):
426
+ flat_tensor = flat_groups_at_rank_i[group_id]
427
+ start_offset = self.offset - self.flat_groups_offset[group_id]
428
+ end_offset = min(end_idx, self.flat_groups_offset[group_id + 1]) - self.flat_groups_offset[group_id]
429
+ pad_flat_param_chunks.append(flat_tensor[start_offset:end_offset])
430
+
431
+ # collect weights from all ranks
432
+ pad_flat_param = torch.cat(pad_flat_param_chunks, dim=0)
433
+ param = pad_flat_param[:self.shape.numel()].view(self.shape).contiguous()
434
+ return param
435
+
436
+
437
+ def _zero3_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states):
438
+ param_shapes = zero_model_states[0].param_shapes
439
+ avail_numel = sum([flat_group.numel() for flat_group in fp32_flat_groups[0]]) * world_size
440
+
441
+ # Reconstruction protocol: For zero3 we need to zip the partitions together at boundary of each
442
+ # param, re-consolidating each param, while dealing with padding if any
443
+
444
+ # merge list of dicts, preserving order
445
+ param_shapes = {k: v for d in param_shapes for k, v in d.items()}
446
+
447
+ if debug:
448
+ for i in range(world_size):
449
+ print(f"{FP32_FLAT_GROUPS}[{i}].shape={fp32_flat_groups[i].shape}")
450
+
451
+ wanted_params = len(param_shapes)
452
+ wanted_numel = sum(shape.numel() for shape in param_shapes.values())
453
+ # not asserting if there is a mismatch due to possible padding
454
+ avail_numel = fp32_flat_groups[0].numel() * world_size
455
+ print(f"Trainable params: Have {avail_numel} numels to process.")
456
+ print(f"Trainable params: Need {wanted_numel} numels in {wanted_params} params.")
457
+
458
+ # params
459
+ # XXX: for huge models that can't fit into the host's RAM we will have to recode this to support
460
+ # out-of-core computing solution
461
+ offset = 0
462
+ total_numel = 0
463
+ total_params = 0
464
+ flat_groups_offset = [0] + list(np.cumsum([flat_tensor.numel() for flat_tensor in fp32_flat_groups[0]]))
465
+ for name, shape in tqdm(param_shapes.items(), desc='Gathering sharded weights'):
466
+ unpartitioned_numel = shape.numel()
467
+ total_numel += unpartitioned_numel
468
+ total_params += 1
469
+ partitioned_numel, partitioned_padding_numel = zero3_partitioned_param_info(unpartitioned_numel, world_size)
470
+
471
+ if debug:
472
+ print(
473
+ f"Trainable params: {total_params} {name} full shape: {shape} partition0 numel={partitioned_numel} partitioned_padding_numel={partitioned_padding_numel}"
474
+ )
475
+
476
+ # memory efficient tensor
477
+ tensor = GatheredTensor(fp32_flat_groups, flat_groups_offset, offset, partitioned_numel, shape)
478
+ state_dict[name] = tensor
479
+ offset += partitioned_numel
480
+
481
+ offset *= world_size
482
+
483
+ # Sanity check
484
+ if offset != avail_numel:
485
+ raise ValueError(f"consumed {offset} numels out of {avail_numel} - something is wrong")
486
+
487
+ print(f"Reconstructed Trainable fp32 state dict with {total_params} params {total_numel} elements")
488
+
489
+
490
+ def _get_fp32_state_dict_from_zero3_checkpoint(world_size, fp32_flat_groups, zero_model_states,
491
+ exclude_frozen_parameters):
492
+ state_dict = OrderedDict()
493
+
494
+ # buffers
495
+ buffers = zero_model_states[0].buffers
496
+ state_dict.update(buffers)
497
+ if debug:
498
+ print(f"added {len(buffers)} buffers")
499
+
500
+ if not exclude_frozen_parameters:
501
+ _zero3_merge_frozen_params(state_dict, world_size, zero_model_states)
502
+
503
+ _zero3_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states)
504
+
505
+ # recover shared parameters
506
+ for pair in zero_model_states[0].shared_params:
507
+ if pair[1] in state_dict:
508
+ state_dict[pair[0]] = state_dict[pair[1]]
509
+
510
+ return state_dict
511
+
512
+
513
+ def to_torch_tensor(state_dict, return_empty_tensor=False):
514
+ """
515
+ Convert state_dict of GatheredTensor to torch tensor
516
+ """
517
+ torch_state_dict = {}
518
+ converted_tensors = {}
519
+ for name, tensor in state_dict.items():
520
+ tensor_id = id(tensor)
521
+ if tensor_id in converted_tensors: # shared tensors
522
+ shared_tensor = torch_state_dict[converted_tensors[tensor_id]]
523
+ torch_state_dict[name] = shared_tensor
524
+ else:
525
+ converted_tensors[tensor_id] = name
526
+ if return_empty_tensor:
527
+ torch_state_dict[name] = torch.empty(tensor.shape, dtype=tensor.dtype)
528
+ else:
529
+ torch_state_dict[name] = tensor.contiguous()
530
+ return torch_state_dict
531
+
532
+
533
+ def get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir,
534
+ tag=None,
535
+ exclude_frozen_parameters=False,
536
+ lazy_mode=False):
537
+ """
538
+ Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated state_dict that can be loaded with
539
+ ``load_state_dict()`` and used for training without DeepSpeed or shared with others, for example
540
+ via a model hub.
541
+
542
+ Args:
543
+ - ``checkpoint_dir``: path to the desired checkpoint folder
544
+ - ``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``
545
+ - ``exclude_frozen_parameters``: exclude frozen parameters
546
+ - ``lazy_mode``: get state_dict in lazy mode. It returns a dict of pesduo tensor instead of torch tensor, which is more memory efficient.
547
+ Convert the pesduo tensor to torch tensor by ``.contiguous()``
548
+
549
+ Returns:
550
+ - pytorch ``state_dict``
551
+
552
+ A typical usage might be ::
553
+
554
+ from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint
555
+ # do the training and checkpoint saving
556
+ state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir) # already on cpu
557
+ model = model.cpu() # move to cpu
558
+ model.load_state_dict(state_dict)
559
+ # submit to model hub or save the model to share with others
560
+
561
+ In this example the ``model`` will no longer be usable in the deepspeed context of the same
562
+ application. i.e. you will need to re-initialize the deepspeed engine, since
563
+ ``model.load_state_dict(state_dict)`` will remove all the deepspeed magic from it.
564
+
565
+ If you want it all done for you, use ``load_state_dict_from_zero_checkpoint`` instead.
566
+
567
+ Note: the above usage may not work if your application doesn't have sufficient free CPU memory.
568
+ You may need to use the offline approach using the ``zero_to_fp32.py`` script that is saved with
569
+ the checkpoint. Or you can load state_dict in lazy mode ::
570
+
571
+ from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint
572
+ state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, lazy_mode=True) # not on cpu
573
+ for name, lazy_tensor in state_dict.item():
574
+ tensor = lazy_tensor.contiguous() # to cpu
575
+ print(name, tensor)
576
+ # del tensor to release memory if it no longer in use
577
+ """
578
+ if tag is None:
579
+ latest_path = os.path.join(checkpoint_dir, 'latest')
580
+ if os.path.isfile(latest_path):
581
+ with open(latest_path, 'r') as fd:
582
+ tag = fd.read().strip()
583
+ else:
584
+ raise ValueError(f"Unable to find 'latest' file at {latest_path}")
585
+
586
+ ds_checkpoint_dir = os.path.join(checkpoint_dir, tag)
587
+
588
+ if not os.path.isdir(ds_checkpoint_dir):
589
+ raise FileNotFoundError(f"Directory '{ds_checkpoint_dir}' doesn't exist")
590
+
591
+ state_dict = _get_fp32_state_dict_from_zero_checkpoint(ds_checkpoint_dir, exclude_frozen_parameters)
592
+ if lazy_mode:
593
+ return state_dict
594
+ else:
595
+ return to_torch_tensor(state_dict)
596
+
597
+
598
+ def convert_zero_checkpoint_to_fp32_state_dict(checkpoint_dir,
599
+ output_dir,
600
+ max_shard_size="5GB",
601
+ safe_serialization=False,
602
+ tag=None,
603
+ exclude_frozen_parameters=False):
604
+ """
605
+ Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated ``state_dict`` file that can be
606
+ loaded with ``torch.load(file)`` + ``load_state_dict()`` and used for training without DeepSpeed.
607
+
608
+ Args:
609
+ - ``checkpoint_dir``: path to the desired checkpoint folder. (one that contains the tag-folder, like ``global_step14``)
610
+ - ``output_dir``: directory to the pytorch fp32 state_dict output files
611
+ - ``max_shard_size``: the maximum size for a checkpoint before being sharded, default value is 5GB
612
+ - ``safe_serialization``: whether to save the model using `safetensors` or the traditional PyTorch way (that uses `pickle`).
613
+ - ``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``
614
+ - ``exclude_frozen_parameters``: exclude frozen parameters
615
+ """
616
+
617
+ # Dependency pre-check
618
+ if safe_serialization:
619
+ try:
620
+ from safetensors.torch import save_file
621
+ except ImportError:
622
+ print('If you want to use `safe_serialization`, please `pip install safetensors`')
623
+ raise
624
+ if max_shard_size is not None:
625
+ try:
626
+ from huggingface_hub import split_torch_state_dict_into_shards
627
+ except ImportError:
628
+ print('If you want to use `max_shard_size`, please `pip install huggingface_hub`')
629
+ raise
630
+
631
+ # Convert zero checkpoint to state_dict
632
+ state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir,
633
+ tag,
634
+ exclude_frozen_parameters,
635
+ lazy_mode=True)
636
+
637
+ # Shard the model if it is too big.
638
+ weights_name = "model.safetensors" if safe_serialization else "pytorch_model.bin"
639
+ if max_shard_size is not None:
640
+ filename_pattern = weights_name.replace(".bin", "{suffix}.bin").replace(".safetensors", "{suffix}.safetensors")
641
+ # an memory-efficient approach for sharding
642
+ empty_state_dict = to_torch_tensor(state_dict, return_empty_tensor=True)
643
+ state_dict_split = split_torch_state_dict_into_shards(empty_state_dict,
644
+ filename_pattern=filename_pattern,
645
+ max_shard_size=max_shard_size)
646
+ else:
647
+ from collections import namedtuple
648
+ StateDictSplit = namedtuple("StateDictSplit", ["is_sharded", "filename_to_tensors"])
649
+ state_dict_split = StateDictSplit(is_sharded=False,
650
+ filename_to_tensors={weights_name: list(state_dict.keys())})
651
+
652
+ # Save the model by shard
653
+ os.makedirs(output_dir, exist_ok=True)
654
+ filename_to_tensors = state_dict_split.filename_to_tensors.items()
655
+ for shard_file, tensors in tqdm(filename_to_tensors, desc="Saving checkpoint shards"):
656
+ shard_state_dict = {tensor_name: state_dict[tensor_name] for tensor_name in tensors}
657
+ shard_state_dict = to_torch_tensor(shard_state_dict)
658
+ output_path = os.path.join(output_dir, shard_file)
659
+ if safe_serialization:
660
+ save_file(shard_state_dict, output_path, metadata={"format": "pt"})
661
+ else:
662
+ torch.save(shard_state_dict, output_path)
663
+ # release the memory of current shard
664
+ for tensor_name in list(shard_state_dict.keys()):
665
+ del state_dict[tensor_name]
666
+ del shard_state_dict[tensor_name]
667
+ del shard_state_dict
668
+ gc.collect()
669
+
670
+ # Save index if sharded
671
+ if state_dict_split.is_sharded:
672
+ index = {
673
+ "metadata": state_dict_split.metadata,
674
+ "weight_map": state_dict_split.tensor_to_filename,
675
+ }
676
+ save_index_file = "model.safetensors.index.json" if safe_serialization else "pytorch_model.bin.index.json"
677
+ save_index_file = os.path.join(output_dir, save_index_file)
678
+ with open(save_index_file, "w", encoding="utf-8") as f:
679
+ content = json.dumps(index, indent=2, sort_keys=True) + "\n"
680
+ f.write(content)
681
+
682
+
683
+ def load_state_dict_from_zero_checkpoint(model, checkpoint_dir, tag=None):
684
+ """
685
+ 1. Put the provided model to cpu
686
+ 2. Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated ``state_dict``
687
+ 3. Load it into the provided model
688
+
689
+ Args:
690
+ - ``model``: the model object to update
691
+ - ``checkpoint_dir``: path to the desired checkpoint folder. (one that contains the tag-folder, like ``global_step14``)
692
+ - ``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``
693
+
694
+ Returns:
695
+ - ``model`: modified model
696
+
697
+ Make sure you have plenty of CPU memory available before you call this function. If you don't
698
+ have enough use the ``zero_to_fp32.py`` utility to do the conversion. You will find it
699
+ conveniently placed for you in the checkpoint folder.
700
+
701
+ A typical usage might be ::
702
+
703
+ from deepspeed.utils.zero_to_fp32 import load_state_dict_from_zero_checkpoint
704
+ model = load_state_dict_from_zero_checkpoint(trainer.model, checkpoint_dir)
705
+ # submit to model hub or save the model to share with others
706
+
707
+ Note, that once this was run, the ``model`` will no longer be usable in the deepspeed context
708
+ of the same application. i.e. you will need to re-initialize the deepspeed engine, since
709
+ ``model.load_state_dict(state_dict)`` will remove all the deepspeed magic from it.
710
+
711
+ """
712
+ logger.info("Extracting fp32 weights")
713
+ state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, tag)
714
+
715
+ logger.info("Overwriting model with fp32 weights")
716
+ model = model.cpu()
717
+ model.load_state_dict(state_dict, strict=False)
718
+
719
+ return model
720
+
721
+
722
+ if __name__ == "__main__":
723
+ parser = argparse.ArgumentParser()
724
+ parser.add_argument("checkpoint_dir",
725
+ type=str,
726
+ help="path to the desired checkpoint folder, e.g., path/checkpoint-12")
727
+ parser.add_argument("output_dir",
728
+ type=str,
729
+ help="directory to the pytorch fp32 state_dict output files"
730
+ "(e.g. path/checkpoint-12-output/)")
731
+ parser.add_argument(
732
+ "--max_shard_size",
733
+ type=str,
734
+ default="5GB",
735
+ help="The maximum size for a checkpoint before being sharded. Checkpoints shard will then be each of size"
736
+ "lower than this size. If expressed as a string, needs to be digits followed by a unit (like `5MB`"
737
+ "We default it to 5GB in order for models to be able to run easily on free-tier google colab instances"
738
+ "without CPU OOM issues.")
739
+ parser.add_argument(
740
+ "--safe_serialization",
741
+ default=False,
742
+ action='store_true',
743
+ help="Whether to save the model using `safetensors` or the traditional PyTorch way (that uses `pickle`).")
744
+ parser.add_argument("-t",
745
+ "--tag",
746
+ type=str,
747
+ default=None,
748
+ help="checkpoint tag used as a unique identifier for checkpoint. e.g., global_step1")
749
+ parser.add_argument("--exclude_frozen_parameters", action='store_true', help="exclude frozen parameters")
750
+ parser.add_argument("-d", "--debug", action='store_true', help="enable debug")
751
+ args = parser.parse_args()
752
+
753
+ debug = args.debug
754
+
755
+ convert_zero_checkpoint_to_fp32_state_dict(args.checkpoint_dir,
756
+ args.output_dir,
757
+ max_shard_size=args.max_shard_size,
758
+ safe_serialization=args.safe_serialization,
759
+ tag=args.tag,
760
+ exclude_frozen_parameters=args.exclude_frozen_parameters)
openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1640/config.json ADDED
@@ -0,0 +1,39 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "architectures": [
3
+ "LlamaForCausalLM"
4
+ ],
5
+ "attention_bias": false,
6
+ "attention_dropout": 0.0,
7
+ "bos_token_id": 128000,
8
+ "dtype": "bfloat16",
9
+ "eos_token_id": [
10
+ 128001,
11
+ 128008,
12
+ 128009
13
+ ],
14
+ "head_dim": 64,
15
+ "hidden_act": "silu",
16
+ "hidden_size": 2048,
17
+ "initializer_range": 0.02,
18
+ "intermediate_size": 8192,
19
+ "max_position_embeddings": 131072,
20
+ "mlp_bias": false,
21
+ "model_type": "llama",
22
+ "num_attention_heads": 32,
23
+ "num_hidden_layers": 16,
24
+ "num_key_value_heads": 8,
25
+ "pretraining_tp": 1,
26
+ "rms_norm_eps": 1e-05,
27
+ "rope_scaling": {
28
+ "factor": 32.0,
29
+ "high_freq_factor": 4.0,
30
+ "low_freq_factor": 1.0,
31
+ "original_max_position_embeddings": 8192,
32
+ "rope_type": "llama3"
33
+ },
34
+ "rope_theta": 500000.0,
35
+ "tie_word_embeddings": true,
36
+ "transformers_version": "4.57.1",
37
+ "use_cache": false,
38
+ "vocab_size": 128256
39
+ }
openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1640/generation_config.json ADDED
@@ -0,0 +1,12 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "bos_token_id": 128000,
3
+ "do_sample": true,
4
+ "eos_token_id": [
5
+ 128001,
6
+ 128008,
7
+ 128009
8
+ ],
9
+ "temperature": 0.6,
10
+ "top_p": 0.9,
11
+ "transformers_version": "4.57.1"
12
+ }
openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1640/latest ADDED
@@ -0,0 +1 @@
 
 
1
+ global_step1640
openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1640/model.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:8257b1ffe27bccf980bead367bedf21a6cbdab1f058457747062cf9c26601b19
3
+ size 2996982344
openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1640/scheduler.pt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:871123a3f5e5af48623dacbf8d53b1f1d8ad18307a63bd7259d3c594eb7e872a
3
+ size 1465
openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1640/trainer_state.json ADDED
@@ -0,0 +1,1182 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "best_global_step": null,
3
+ "best_metric": null,
4
+ "best_model_checkpoint": null,
5
+ "epoch": 1.9981718464351006,
6
+ "eval_steps": 500,
7
+ "global_step": 1640,
8
+ "is_hyper_param_search": false,
9
+ "is_local_process_zero": true,
10
+ "is_world_process_zero": true,
11
+ "log_history": [
12
+ {
13
+ "epoch": 0.01218769043266301,
14
+ "grad_norm": 0.36486390233039856,
15
+ "learning_rate": 5e-05,
16
+ "loss": 1.2213,
17
+ "step": 10
18
+ },
19
+ {
20
+ "epoch": 0.02437538086532602,
21
+ "grad_norm": 0.2953564524650574,
22
+ "learning_rate": 5e-05,
23
+ "loss": 1.0223,
24
+ "step": 20
25
+ },
26
+ {
27
+ "epoch": 0.03656307129798903,
28
+ "grad_norm": 0.2766735851764679,
29
+ "learning_rate": 5e-05,
30
+ "loss": 0.9589,
31
+ "step": 30
32
+ },
33
+ {
34
+ "epoch": 0.04875076173065204,
35
+ "grad_norm": 0.32770997285842896,
36
+ "learning_rate": 5e-05,
37
+ "loss": 0.9081,
38
+ "step": 40
39
+ },
40
+ {
41
+ "epoch": 0.06093845216331505,
42
+ "grad_norm": 0.3052276372909546,
43
+ "learning_rate": 5e-05,
44
+ "loss": 0.8752,
45
+ "step": 50
46
+ },
47
+ {
48
+ "epoch": 0.07312614259597806,
49
+ "grad_norm": 0.240146666765213,
50
+ "learning_rate": 5e-05,
51
+ "loss": 0.873,
52
+ "step": 60
53
+ },
54
+ {
55
+ "epoch": 0.08531383302864107,
56
+ "grad_norm": 0.23872187733650208,
57
+ "learning_rate": 5e-05,
58
+ "loss": 0.8587,
59
+ "step": 70
60
+ },
61
+ {
62
+ "epoch": 0.09750152346130408,
63
+ "grad_norm": 0.32638487219810486,
64
+ "learning_rate": 5e-05,
65
+ "loss": 0.8656,
66
+ "step": 80
67
+ },
68
+ {
69
+ "epoch": 0.10968921389396709,
70
+ "grad_norm": 0.25019028782844543,
71
+ "learning_rate": 5e-05,
72
+ "loss": 0.8685,
73
+ "step": 90
74
+ },
75
+ {
76
+ "epoch": 0.1218769043266301,
77
+ "grad_norm": 0.2507675886154175,
78
+ "learning_rate": 5e-05,
79
+ "loss": 0.8662,
80
+ "step": 100
81
+ },
82
+ {
83
+ "epoch": 0.1340645947592931,
84
+ "grad_norm": 0.27966755628585815,
85
+ "learning_rate": 5e-05,
86
+ "loss": 0.8428,
87
+ "step": 110
88
+ },
89
+ {
90
+ "epoch": 0.14625228519195613,
91
+ "grad_norm": 0.2622673213481903,
92
+ "learning_rate": 5e-05,
93
+ "loss": 0.8416,
94
+ "step": 120
95
+ },
96
+ {
97
+ "epoch": 0.15843997562461914,
98
+ "grad_norm": 0.3003966510295868,
99
+ "learning_rate": 5e-05,
100
+ "loss": 0.8528,
101
+ "step": 130
102
+ },
103
+ {
104
+ "epoch": 0.17062766605728213,
105
+ "grad_norm": 0.32308319211006165,
106
+ "learning_rate": 5e-05,
107
+ "loss": 0.8357,
108
+ "step": 140
109
+ },
110
+ {
111
+ "epoch": 0.18281535648994515,
112
+ "grad_norm": 0.27279096841812134,
113
+ "learning_rate": 5e-05,
114
+ "loss": 0.8432,
115
+ "step": 150
116
+ },
117
+ {
118
+ "epoch": 0.19500304692260817,
119
+ "grad_norm": 0.27976104617118835,
120
+ "learning_rate": 5e-05,
121
+ "loss": 0.8119,
122
+ "step": 160
123
+ },
124
+ {
125
+ "epoch": 0.2071907373552712,
126
+ "grad_norm": 0.290866881608963,
127
+ "learning_rate": 5e-05,
128
+ "loss": 0.8096,
129
+ "step": 170
130
+ },
131
+ {
132
+ "epoch": 0.21937842778793418,
133
+ "grad_norm": 0.29430949687957764,
134
+ "learning_rate": 5e-05,
135
+ "loss": 0.819,
136
+ "step": 180
137
+ },
138
+ {
139
+ "epoch": 0.2315661182205972,
140
+ "grad_norm": 0.27714136242866516,
141
+ "learning_rate": 5e-05,
142
+ "loss": 0.8187,
143
+ "step": 190
144
+ },
145
+ {
146
+ "epoch": 0.2437538086532602,
147
+ "grad_norm": 0.2997029423713684,
148
+ "learning_rate": 5e-05,
149
+ "loss": 0.808,
150
+ "step": 200
151
+ },
152
+ {
153
+ "epoch": 0.25594149908592323,
154
+ "grad_norm": 0.2539573609828949,
155
+ "learning_rate": 5e-05,
156
+ "loss": 0.7943,
157
+ "step": 210
158
+ },
159
+ {
160
+ "epoch": 0.2681291895185862,
161
+ "grad_norm": 0.23294171690940857,
162
+ "learning_rate": 5e-05,
163
+ "loss": 0.8188,
164
+ "step": 220
165
+ },
166
+ {
167
+ "epoch": 0.28031687995124926,
168
+ "grad_norm": 0.25016292929649353,
169
+ "learning_rate": 5e-05,
170
+ "loss": 0.8035,
171
+ "step": 230
172
+ },
173
+ {
174
+ "epoch": 0.29250457038391225,
175
+ "grad_norm": 0.37125322222709656,
176
+ "learning_rate": 5e-05,
177
+ "loss": 0.785,
178
+ "step": 240
179
+ },
180
+ {
181
+ "epoch": 0.30469226081657524,
182
+ "grad_norm": 0.21520452201366425,
183
+ "learning_rate": 5e-05,
184
+ "loss": 0.8088,
185
+ "step": 250
186
+ },
187
+ {
188
+ "epoch": 0.3168799512492383,
189
+ "grad_norm": 0.20990602672100067,
190
+ "learning_rate": 5e-05,
191
+ "loss": 0.7813,
192
+ "step": 260
193
+ },
194
+ {
195
+ "epoch": 0.3290676416819013,
196
+ "grad_norm": 0.24193961918354034,
197
+ "learning_rate": 5e-05,
198
+ "loss": 0.7979,
199
+ "step": 270
200
+ },
201
+ {
202
+ "epoch": 0.34125533211456427,
203
+ "grad_norm": 0.22766050696372986,
204
+ "learning_rate": 5e-05,
205
+ "loss": 0.785,
206
+ "step": 280
207
+ },
208
+ {
209
+ "epoch": 0.3534430225472273,
210
+ "grad_norm": 0.2662702202796936,
211
+ "learning_rate": 5e-05,
212
+ "loss": 0.7766,
213
+ "step": 290
214
+ },
215
+ {
216
+ "epoch": 0.3656307129798903,
217
+ "grad_norm": 0.23020495474338531,
218
+ "learning_rate": 5e-05,
219
+ "loss": 0.7687,
220
+ "step": 300
221
+ },
222
+ {
223
+ "epoch": 0.3778184034125533,
224
+ "grad_norm": 0.21770520508289337,
225
+ "learning_rate": 5e-05,
226
+ "loss": 0.7815,
227
+ "step": 310
228
+ },
229
+ {
230
+ "epoch": 0.39000609384521634,
231
+ "grad_norm": 0.22677403688430786,
232
+ "learning_rate": 5e-05,
233
+ "loss": 0.7999,
234
+ "step": 320
235
+ },
236
+ {
237
+ "epoch": 0.40219378427787933,
238
+ "grad_norm": 0.24891860783100128,
239
+ "learning_rate": 5e-05,
240
+ "loss": 0.7946,
241
+ "step": 330
242
+ },
243
+ {
244
+ "epoch": 0.4143814747105424,
245
+ "grad_norm": 0.2506854832172394,
246
+ "learning_rate": 5e-05,
247
+ "loss": 0.7889,
248
+ "step": 340
249
+ },
250
+ {
251
+ "epoch": 0.42656916514320536,
252
+ "grad_norm": 0.24388180673122406,
253
+ "learning_rate": 5e-05,
254
+ "loss": 0.79,
255
+ "step": 350
256
+ },
257
+ {
258
+ "epoch": 0.43875685557586835,
259
+ "grad_norm": 0.22329650819301605,
260
+ "learning_rate": 5e-05,
261
+ "loss": 0.7632,
262
+ "step": 360
263
+ },
264
+ {
265
+ "epoch": 0.4509445460085314,
266
+ "grad_norm": 0.2240927815437317,
267
+ "learning_rate": 5e-05,
268
+ "loss": 0.8042,
269
+ "step": 370
270
+ },
271
+ {
272
+ "epoch": 0.4631322364411944,
273
+ "grad_norm": 0.2578203082084656,
274
+ "learning_rate": 5e-05,
275
+ "loss": 0.7667,
276
+ "step": 380
277
+ },
278
+ {
279
+ "epoch": 0.4753199268738574,
280
+ "grad_norm": 0.25349217653274536,
281
+ "learning_rate": 5e-05,
282
+ "loss": 0.7911,
283
+ "step": 390
284
+ },
285
+ {
286
+ "epoch": 0.4875076173065204,
287
+ "grad_norm": 0.254446804523468,
288
+ "learning_rate": 5e-05,
289
+ "loss": 0.7909,
290
+ "step": 400
291
+ },
292
+ {
293
+ "epoch": 0.4996953077391834,
294
+ "grad_norm": 0.24253171682357788,
295
+ "learning_rate": 5e-05,
296
+ "loss": 0.7896,
297
+ "step": 410
298
+ },
299
+ {
300
+ "epoch": 0.5118829981718465,
301
+ "grad_norm": 0.27901458740234375,
302
+ "learning_rate": 5e-05,
303
+ "loss": 0.7828,
304
+ "step": 420
305
+ },
306
+ {
307
+ "epoch": 0.5240706886045094,
308
+ "grad_norm": 0.31884434819221497,
309
+ "learning_rate": 5e-05,
310
+ "loss": 0.7919,
311
+ "step": 430
312
+ },
313
+ {
314
+ "epoch": 0.5362583790371724,
315
+ "grad_norm": 0.22920799255371094,
316
+ "learning_rate": 5e-05,
317
+ "loss": 0.7777,
318
+ "step": 440
319
+ },
320
+ {
321
+ "epoch": 0.5484460694698354,
322
+ "grad_norm": 0.22606618702411652,
323
+ "learning_rate": 5e-05,
324
+ "loss": 0.7661,
325
+ "step": 450
326
+ },
327
+ {
328
+ "epoch": 0.5606337599024985,
329
+ "grad_norm": 0.21604052186012268,
330
+ "learning_rate": 5e-05,
331
+ "loss": 0.769,
332
+ "step": 460
333
+ },
334
+ {
335
+ "epoch": 0.5728214503351615,
336
+ "grad_norm": 0.2296101152896881,
337
+ "learning_rate": 5e-05,
338
+ "loss": 0.7626,
339
+ "step": 470
340
+ },
341
+ {
342
+ "epoch": 0.5850091407678245,
343
+ "grad_norm": 0.24162150919437408,
344
+ "learning_rate": 5e-05,
345
+ "loss": 0.7831,
346
+ "step": 480
347
+ },
348
+ {
349
+ "epoch": 0.5971968312004875,
350
+ "grad_norm": 0.22364799678325653,
351
+ "learning_rate": 5e-05,
352
+ "loss": 0.7737,
353
+ "step": 490
354
+ },
355
+ {
356
+ "epoch": 0.6093845216331505,
357
+ "grad_norm": 0.23953403532505035,
358
+ "learning_rate": 5e-05,
359
+ "loss": 0.787,
360
+ "step": 500
361
+ },
362
+ {
363
+ "epoch": 0.6215722120658135,
364
+ "grad_norm": 0.22349311411380768,
365
+ "learning_rate": 5e-05,
366
+ "loss": 0.7763,
367
+ "step": 510
368
+ },
369
+ {
370
+ "epoch": 0.6337599024984766,
371
+ "grad_norm": 0.21996377408504486,
372
+ "learning_rate": 5e-05,
373
+ "loss": 0.7596,
374
+ "step": 520
375
+ },
376
+ {
377
+ "epoch": 0.6459475929311396,
378
+ "grad_norm": 0.245538130402565,
379
+ "learning_rate": 5e-05,
380
+ "loss": 0.7785,
381
+ "step": 530
382
+ },
383
+ {
384
+ "epoch": 0.6581352833638026,
385
+ "grad_norm": 0.23442520201206207,
386
+ "learning_rate": 5e-05,
387
+ "loss": 0.7633,
388
+ "step": 540
389
+ },
390
+ {
391
+ "epoch": 0.6703229737964655,
392
+ "grad_norm": 0.22274106740951538,
393
+ "learning_rate": 5e-05,
394
+ "loss": 0.7756,
395
+ "step": 550
396
+ },
397
+ {
398
+ "epoch": 0.6825106642291285,
399
+ "grad_norm": 0.21717530488967896,
400
+ "learning_rate": 5e-05,
401
+ "loss": 0.7726,
402
+ "step": 560
403
+ },
404
+ {
405
+ "epoch": 0.6946983546617916,
406
+ "grad_norm": 0.2465955913066864,
407
+ "learning_rate": 5e-05,
408
+ "loss": 0.7476,
409
+ "step": 570
410
+ },
411
+ {
412
+ "epoch": 0.7068860450944546,
413
+ "grad_norm": 0.2318575084209442,
414
+ "learning_rate": 5e-05,
415
+ "loss": 0.7584,
416
+ "step": 580
417
+ },
418
+ {
419
+ "epoch": 0.7190737355271176,
420
+ "grad_norm": 0.25114378333091736,
421
+ "learning_rate": 5e-05,
422
+ "loss": 0.7625,
423
+ "step": 590
424
+ },
425
+ {
426
+ "epoch": 0.7312614259597806,
427
+ "grad_norm": 0.19145755469799042,
428
+ "learning_rate": 5e-05,
429
+ "loss": 0.7577,
430
+ "step": 600
431
+ },
432
+ {
433
+ "epoch": 0.7434491163924436,
434
+ "grad_norm": 0.19129066169261932,
435
+ "learning_rate": 5e-05,
436
+ "loss": 0.7772,
437
+ "step": 610
438
+ },
439
+ {
440
+ "epoch": 0.7556368068251066,
441
+ "grad_norm": 0.2014475166797638,
442
+ "learning_rate": 5e-05,
443
+ "loss": 0.7601,
444
+ "step": 620
445
+ },
446
+ {
447
+ "epoch": 0.7678244972577697,
448
+ "grad_norm": 0.21480509638786316,
449
+ "learning_rate": 5e-05,
450
+ "loss": 0.7665,
451
+ "step": 630
452
+ },
453
+ {
454
+ "epoch": 0.7800121876904327,
455
+ "grad_norm": 0.2633228302001953,
456
+ "learning_rate": 5e-05,
457
+ "loss": 0.7565,
458
+ "step": 640
459
+ },
460
+ {
461
+ "epoch": 0.7921998781230957,
462
+ "grad_norm": 0.24633203446865082,
463
+ "learning_rate": 5e-05,
464
+ "loss": 0.7602,
465
+ "step": 650
466
+ },
467
+ {
468
+ "epoch": 0.8043875685557587,
469
+ "grad_norm": 0.19770534336566925,
470
+ "learning_rate": 5e-05,
471
+ "loss": 0.7685,
472
+ "step": 660
473
+ },
474
+ {
475
+ "epoch": 0.8165752589884216,
476
+ "grad_norm": 0.2267930805683136,
477
+ "learning_rate": 5e-05,
478
+ "loss": 0.7516,
479
+ "step": 670
480
+ },
481
+ {
482
+ "epoch": 0.8287629494210847,
483
+ "grad_norm": 0.238111674785614,
484
+ "learning_rate": 5e-05,
485
+ "loss": 0.7487,
486
+ "step": 680
487
+ },
488
+ {
489
+ "epoch": 0.8409506398537477,
490
+ "grad_norm": 0.19528904557228088,
491
+ "learning_rate": 5e-05,
492
+ "loss": 0.7467,
493
+ "step": 690
494
+ },
495
+ {
496
+ "epoch": 0.8531383302864107,
497
+ "grad_norm": 0.21305914223194122,
498
+ "learning_rate": 5e-05,
499
+ "loss": 0.7527,
500
+ "step": 700
501
+ },
502
+ {
503
+ "epoch": 0.8653260207190737,
504
+ "grad_norm": 0.2205205112695694,
505
+ "learning_rate": 5e-05,
506
+ "loss": 0.7639,
507
+ "step": 710
508
+ },
509
+ {
510
+ "epoch": 0.8775137111517367,
511
+ "grad_norm": 0.22429127991199493,
512
+ "learning_rate": 5e-05,
513
+ "loss": 0.7744,
514
+ "step": 720
515
+ },
516
+ {
517
+ "epoch": 0.8897014015843998,
518
+ "grad_norm": 0.23755772411823273,
519
+ "learning_rate": 5e-05,
520
+ "loss": 0.7446,
521
+ "step": 730
522
+ },
523
+ {
524
+ "epoch": 0.9018890920170628,
525
+ "grad_norm": 0.17592693865299225,
526
+ "learning_rate": 5e-05,
527
+ "loss": 0.7442,
528
+ "step": 740
529
+ },
530
+ {
531
+ "epoch": 0.9140767824497258,
532
+ "grad_norm": 0.2266898900270462,
533
+ "learning_rate": 5e-05,
534
+ "loss": 0.7326,
535
+ "step": 750
536
+ },
537
+ {
538
+ "epoch": 0.9262644728823888,
539
+ "grad_norm": 0.23541569709777832,
540
+ "learning_rate": 5e-05,
541
+ "loss": 0.7338,
542
+ "step": 760
543
+ },
544
+ {
545
+ "epoch": 0.9384521633150518,
546
+ "grad_norm": 0.22233888506889343,
547
+ "learning_rate": 5e-05,
548
+ "loss": 0.7385,
549
+ "step": 770
550
+ },
551
+ {
552
+ "epoch": 0.9506398537477148,
553
+ "grad_norm": 0.23714767396450043,
554
+ "learning_rate": 5e-05,
555
+ "loss": 0.7306,
556
+ "step": 780
557
+ },
558
+ {
559
+ "epoch": 0.9628275441803779,
560
+ "grad_norm": 0.21191060543060303,
561
+ "learning_rate": 5e-05,
562
+ "loss": 0.736,
563
+ "step": 790
564
+ },
565
+ {
566
+ "epoch": 0.9750152346130408,
567
+ "grad_norm": 0.22759106755256653,
568
+ "learning_rate": 5e-05,
569
+ "loss": 0.7505,
570
+ "step": 800
571
+ },
572
+ {
573
+ "epoch": 0.9872029250457038,
574
+ "grad_norm": 0.19895903766155243,
575
+ "learning_rate": 5e-05,
576
+ "loss": 0.7536,
577
+ "step": 810
578
+ },
579
+ {
580
+ "epoch": 0.9993906154783668,
581
+ "grad_norm": 0.3061685860157013,
582
+ "learning_rate": 5e-05,
583
+ "loss": 0.7493,
584
+ "step": 820
585
+ },
586
+ {
587
+ "epoch": 1.0109689213893966,
588
+ "grad_norm": 0.18868015706539154,
589
+ "learning_rate": 5e-05,
590
+ "loss": 0.7342,
591
+ "step": 830
592
+ },
593
+ {
594
+ "epoch": 1.0231566118220596,
595
+ "grad_norm": 0.1797628402709961,
596
+ "learning_rate": 5e-05,
597
+ "loss": 0.7192,
598
+ "step": 840
599
+ },
600
+ {
601
+ "epoch": 1.0353443022547226,
602
+ "grad_norm": 0.19511596858501434,
603
+ "learning_rate": 5e-05,
604
+ "loss": 0.7002,
605
+ "step": 850
606
+ },
607
+ {
608
+ "epoch": 1.0475319926873858,
609
+ "grad_norm": 0.2123185396194458,
610
+ "learning_rate": 5e-05,
611
+ "loss": 0.7236,
612
+ "step": 860
613
+ },
614
+ {
615
+ "epoch": 1.0597196831200488,
616
+ "grad_norm": 0.2047470211982727,
617
+ "learning_rate": 5e-05,
618
+ "loss": 0.711,
619
+ "step": 870
620
+ },
621
+ {
622
+ "epoch": 1.0719073735527118,
623
+ "grad_norm": 0.18809868395328522,
624
+ "learning_rate": 5e-05,
625
+ "loss": 0.7253,
626
+ "step": 880
627
+ },
628
+ {
629
+ "epoch": 1.0840950639853748,
630
+ "grad_norm": 0.2207012176513672,
631
+ "learning_rate": 5e-05,
632
+ "loss": 0.6971,
633
+ "step": 890
634
+ },
635
+ {
636
+ "epoch": 1.0962827544180378,
637
+ "grad_norm": 0.1912035495042801,
638
+ "learning_rate": 5e-05,
639
+ "loss": 0.7239,
640
+ "step": 900
641
+ },
642
+ {
643
+ "epoch": 1.1084704448507008,
644
+ "grad_norm": 0.23486360907554626,
645
+ "learning_rate": 5e-05,
646
+ "loss": 0.7335,
647
+ "step": 910
648
+ },
649
+ {
650
+ "epoch": 1.1206581352833638,
651
+ "grad_norm": 0.22189787030220032,
652
+ "learning_rate": 5e-05,
653
+ "loss": 0.7149,
654
+ "step": 920
655
+ },
656
+ {
657
+ "epoch": 1.1328458257160268,
658
+ "grad_norm": 0.25818824768066406,
659
+ "learning_rate": 5e-05,
660
+ "loss": 0.7322,
661
+ "step": 930
662
+ },
663
+ {
664
+ "epoch": 1.1450335161486898,
665
+ "grad_norm": 0.2088775783777237,
666
+ "learning_rate": 5e-05,
667
+ "loss": 0.7382,
668
+ "step": 940
669
+ },
670
+ {
671
+ "epoch": 1.1572212065813527,
672
+ "grad_norm": 0.19445833563804626,
673
+ "learning_rate": 5e-05,
674
+ "loss": 0.7223,
675
+ "step": 950
676
+ },
677
+ {
678
+ "epoch": 1.1694088970140157,
679
+ "grad_norm": 0.21489456295967102,
680
+ "learning_rate": 5e-05,
681
+ "loss": 0.7317,
682
+ "step": 960
683
+ },
684
+ {
685
+ "epoch": 1.181596587446679,
686
+ "grad_norm": 0.24050794541835785,
687
+ "learning_rate": 5e-05,
688
+ "loss": 0.7275,
689
+ "step": 970
690
+ },
691
+ {
692
+ "epoch": 1.193784277879342,
693
+ "grad_norm": 0.18142102658748627,
694
+ "learning_rate": 5e-05,
695
+ "loss": 0.7046,
696
+ "step": 980
697
+ },
698
+ {
699
+ "epoch": 1.205971968312005,
700
+ "grad_norm": 0.4781077802181244,
701
+ "learning_rate": 5e-05,
702
+ "loss": 0.7433,
703
+ "step": 990
704
+ },
705
+ {
706
+ "epoch": 1.218159658744668,
707
+ "grad_norm": 0.20628023147583008,
708
+ "learning_rate": 5e-05,
709
+ "loss": 0.7403,
710
+ "step": 1000
711
+ },
712
+ {
713
+ "epoch": 1.230347349177331,
714
+ "grad_norm": 0.19117140769958496,
715
+ "learning_rate": 5e-05,
716
+ "loss": 0.7308,
717
+ "step": 1010
718
+ },
719
+ {
720
+ "epoch": 1.242535039609994,
721
+ "grad_norm": 0.2107745110988617,
722
+ "learning_rate": 5e-05,
723
+ "loss": 0.7201,
724
+ "step": 1020
725
+ },
726
+ {
727
+ "epoch": 1.2547227300426569,
728
+ "grad_norm": 0.20551154017448425,
729
+ "learning_rate": 5e-05,
730
+ "loss": 0.7038,
731
+ "step": 1030
732
+ },
733
+ {
734
+ "epoch": 1.2669104204753199,
735
+ "grad_norm": 0.26577022671699524,
736
+ "learning_rate": 5e-05,
737
+ "loss": 0.7096,
738
+ "step": 1040
739
+ },
740
+ {
741
+ "epoch": 1.2790981109079829,
742
+ "grad_norm": 0.22094209492206573,
743
+ "learning_rate": 5e-05,
744
+ "loss": 0.718,
745
+ "step": 1050
746
+ },
747
+ {
748
+ "epoch": 1.291285801340646,
749
+ "grad_norm": 0.18224076926708221,
750
+ "learning_rate": 5e-05,
751
+ "loss": 0.7261,
752
+ "step": 1060
753
+ },
754
+ {
755
+ "epoch": 1.303473491773309,
756
+ "grad_norm": 0.226251482963562,
757
+ "learning_rate": 5e-05,
758
+ "loss": 0.7372,
759
+ "step": 1070
760
+ },
761
+ {
762
+ "epoch": 1.315661182205972,
763
+ "grad_norm": 0.20186461508274078,
764
+ "learning_rate": 5e-05,
765
+ "loss": 0.7074,
766
+ "step": 1080
767
+ },
768
+ {
769
+ "epoch": 1.327848872638635,
770
+ "grad_norm": 0.22229762375354767,
771
+ "learning_rate": 5e-05,
772
+ "loss": 0.7069,
773
+ "step": 1090
774
+ },
775
+ {
776
+ "epoch": 1.340036563071298,
777
+ "grad_norm": 0.1776452213525772,
778
+ "learning_rate": 5e-05,
779
+ "loss": 0.714,
780
+ "step": 1100
781
+ },
782
+ {
783
+ "epoch": 1.352224253503961,
784
+ "grad_norm": 0.20680229365825653,
785
+ "learning_rate": 5e-05,
786
+ "loss": 0.7066,
787
+ "step": 1110
788
+ },
789
+ {
790
+ "epoch": 1.364411943936624,
791
+ "grad_norm": 0.21842031180858612,
792
+ "learning_rate": 5e-05,
793
+ "loss": 0.7281,
794
+ "step": 1120
795
+ },
796
+ {
797
+ "epoch": 1.376599634369287,
798
+ "grad_norm": 0.28037774562835693,
799
+ "learning_rate": 5e-05,
800
+ "loss": 0.7149,
801
+ "step": 1130
802
+ },
803
+ {
804
+ "epoch": 1.38878732480195,
805
+ "grad_norm": 0.20962920784950256,
806
+ "learning_rate": 5e-05,
807
+ "loss": 0.6951,
808
+ "step": 1140
809
+ },
810
+ {
811
+ "epoch": 1.400975015234613,
812
+ "grad_norm": 0.18751871585845947,
813
+ "learning_rate": 5e-05,
814
+ "loss": 0.7101,
815
+ "step": 1150
816
+ },
817
+ {
818
+ "epoch": 1.413162705667276,
819
+ "grad_norm": 0.2722964584827423,
820
+ "learning_rate": 5e-05,
821
+ "loss": 0.7205,
822
+ "step": 1160
823
+ },
824
+ {
825
+ "epoch": 1.425350396099939,
826
+ "grad_norm": 0.20932041108608246,
827
+ "learning_rate": 5e-05,
828
+ "loss": 0.7082,
829
+ "step": 1170
830
+ },
831
+ {
832
+ "epoch": 1.437538086532602,
833
+ "grad_norm": 0.17836257815361023,
834
+ "learning_rate": 5e-05,
835
+ "loss": 0.7232,
836
+ "step": 1180
837
+ },
838
+ {
839
+ "epoch": 1.449725776965265,
840
+ "grad_norm": 0.3180703818798065,
841
+ "learning_rate": 5e-05,
842
+ "loss": 0.7245,
843
+ "step": 1190
844
+ },
845
+ {
846
+ "epoch": 1.4619134673979282,
847
+ "grad_norm": 0.201424703001976,
848
+ "learning_rate": 5e-05,
849
+ "loss": 0.7232,
850
+ "step": 1200
851
+ },
852
+ {
853
+ "epoch": 1.4741011578305911,
854
+ "grad_norm": 0.20228999853134155,
855
+ "learning_rate": 5e-05,
856
+ "loss": 0.7169,
857
+ "step": 1210
858
+ },
859
+ {
860
+ "epoch": 1.4862888482632541,
861
+ "grad_norm": 0.17478394508361816,
862
+ "learning_rate": 5e-05,
863
+ "loss": 0.704,
864
+ "step": 1220
865
+ },
866
+ {
867
+ "epoch": 1.4984765386959171,
868
+ "grad_norm": 0.20061351358890533,
869
+ "learning_rate": 5e-05,
870
+ "loss": 0.6792,
871
+ "step": 1230
872
+ },
873
+ {
874
+ "epoch": 1.5106642291285801,
875
+ "grad_norm": 0.1857377588748932,
876
+ "learning_rate": 5e-05,
877
+ "loss": 0.709,
878
+ "step": 1240
879
+ },
880
+ {
881
+ "epoch": 1.522851919561243,
882
+ "grad_norm": 0.19873379170894623,
883
+ "learning_rate": 5e-05,
884
+ "loss": 0.7007,
885
+ "step": 1250
886
+ },
887
+ {
888
+ "epoch": 1.5350396099939063,
889
+ "grad_norm": 0.19807074964046478,
890
+ "learning_rate": 5e-05,
891
+ "loss": 0.7122,
892
+ "step": 1260
893
+ },
894
+ {
895
+ "epoch": 1.5472273004265693,
896
+ "grad_norm": 0.16746796667575836,
897
+ "learning_rate": 5e-05,
898
+ "loss": 0.7005,
899
+ "step": 1270
900
+ },
901
+ {
902
+ "epoch": 1.5594149908592323,
903
+ "grad_norm": 0.2750326991081238,
904
+ "learning_rate": 5e-05,
905
+ "loss": 0.7166,
906
+ "step": 1280
907
+ },
908
+ {
909
+ "epoch": 1.5716026812918953,
910
+ "grad_norm": 0.2776738107204437,
911
+ "learning_rate": 5e-05,
912
+ "loss": 0.7405,
913
+ "step": 1290
914
+ },
915
+ {
916
+ "epoch": 1.5837903717245583,
917
+ "grad_norm": 0.3494517207145691,
918
+ "learning_rate": 5e-05,
919
+ "loss": 0.7192,
920
+ "step": 1300
921
+ },
922
+ {
923
+ "epoch": 1.5959780621572213,
924
+ "grad_norm": 0.25028494000434875,
925
+ "learning_rate": 5e-05,
926
+ "loss": 0.741,
927
+ "step": 1310
928
+ },
929
+ {
930
+ "epoch": 1.6081657525898843,
931
+ "grad_norm": 0.23914143443107605,
932
+ "learning_rate": 5e-05,
933
+ "loss": 0.72,
934
+ "step": 1320
935
+ },
936
+ {
937
+ "epoch": 1.6203534430225472,
938
+ "grad_norm": 0.2052862048149109,
939
+ "learning_rate": 5e-05,
940
+ "loss": 0.7112,
941
+ "step": 1330
942
+ },
943
+ {
944
+ "epoch": 1.6325411334552102,
945
+ "grad_norm": 0.22285936772823334,
946
+ "learning_rate": 5e-05,
947
+ "loss": 0.7165,
948
+ "step": 1340
949
+ },
950
+ {
951
+ "epoch": 1.6447288238878732,
952
+ "grad_norm": 0.20810848474502563,
953
+ "learning_rate": 5e-05,
954
+ "loss": 0.7122,
955
+ "step": 1350
956
+ },
957
+ {
958
+ "epoch": 1.6569165143205362,
959
+ "grad_norm": 0.20422884821891785,
960
+ "learning_rate": 5e-05,
961
+ "loss": 0.7082,
962
+ "step": 1360
963
+ },
964
+ {
965
+ "epoch": 1.6691042047531992,
966
+ "grad_norm": 0.18164314329624176,
967
+ "learning_rate": 5e-05,
968
+ "loss": 0.7131,
969
+ "step": 1370
970
+ },
971
+ {
972
+ "epoch": 1.6812918951858622,
973
+ "grad_norm": 0.19893014430999756,
974
+ "learning_rate": 5e-05,
975
+ "loss": 0.7055,
976
+ "step": 1380
977
+ },
978
+ {
979
+ "epoch": 1.6934795856185252,
980
+ "grad_norm": 0.24865971505641937,
981
+ "learning_rate": 5e-05,
982
+ "loss": 0.6994,
983
+ "step": 1390
984
+ },
985
+ {
986
+ "epoch": 1.7056672760511882,
987
+ "grad_norm": 0.18078123033046722,
988
+ "learning_rate": 5e-05,
989
+ "loss": 0.7162,
990
+ "step": 1400
991
+ },
992
+ {
993
+ "epoch": 1.7178549664838512,
994
+ "grad_norm": 0.1973421722650528,
995
+ "learning_rate": 5e-05,
996
+ "loss": 0.7136,
997
+ "step": 1410
998
+ },
999
+ {
1000
+ "epoch": 1.7300426569165142,
1001
+ "grad_norm": 0.2057911604642868,
1002
+ "learning_rate": 5e-05,
1003
+ "loss": 0.7004,
1004
+ "step": 1420
1005
+ },
1006
+ {
1007
+ "epoch": 1.7422303473491774,
1008
+ "grad_norm": 0.18640495836734772,
1009
+ "learning_rate": 5e-05,
1010
+ "loss": 0.7102,
1011
+ "step": 1430
1012
+ },
1013
+ {
1014
+ "epoch": 1.7544180377818404,
1015
+ "grad_norm": 0.21038417518138885,
1016
+ "learning_rate": 5e-05,
1017
+ "loss": 0.6908,
1018
+ "step": 1440
1019
+ },
1020
+ {
1021
+ "epoch": 1.7666057282145033,
1022
+ "grad_norm": 0.2123890221118927,
1023
+ "learning_rate": 5e-05,
1024
+ "loss": 0.6913,
1025
+ "step": 1450
1026
+ },
1027
+ {
1028
+ "epoch": 1.7787934186471663,
1029
+ "grad_norm": 0.21619991958141327,
1030
+ "learning_rate": 5e-05,
1031
+ "loss": 0.7038,
1032
+ "step": 1460
1033
+ },
1034
+ {
1035
+ "epoch": 1.7909811090798293,
1036
+ "grad_norm": 0.1975220888853073,
1037
+ "learning_rate": 5e-05,
1038
+ "loss": 0.7041,
1039
+ "step": 1470
1040
+ },
1041
+ {
1042
+ "epoch": 1.8031687995124925,
1043
+ "grad_norm": 0.20595207810401917,
1044
+ "learning_rate": 5e-05,
1045
+ "loss": 0.7145,
1046
+ "step": 1480
1047
+ },
1048
+ {
1049
+ "epoch": 1.8153564899451555,
1050
+ "grad_norm": 0.16633334755897522,
1051
+ "learning_rate": 5e-05,
1052
+ "loss": 0.6963,
1053
+ "step": 1490
1054
+ },
1055
+ {
1056
+ "epoch": 1.8275441803778185,
1057
+ "grad_norm": 0.19140620529651642,
1058
+ "learning_rate": 5e-05,
1059
+ "loss": 0.6814,
1060
+ "step": 1500
1061
+ },
1062
+ {
1063
+ "epoch": 1.8397318708104815,
1064
+ "grad_norm": 0.1833905130624771,
1065
+ "learning_rate": 5e-05,
1066
+ "loss": 0.7174,
1067
+ "step": 1510
1068
+ },
1069
+ {
1070
+ "epoch": 1.8519195612431445,
1071
+ "grad_norm": 0.20587337017059326,
1072
+ "learning_rate": 5e-05,
1073
+ "loss": 0.7098,
1074
+ "step": 1520
1075
+ },
1076
+ {
1077
+ "epoch": 1.8641072516758075,
1078
+ "grad_norm": 0.21952909231185913,
1079
+ "learning_rate": 5e-05,
1080
+ "loss": 0.6912,
1081
+ "step": 1530
1082
+ },
1083
+ {
1084
+ "epoch": 1.8762949421084705,
1085
+ "grad_norm": 0.23233620822429657,
1086
+ "learning_rate": 5e-05,
1087
+ "loss": 0.6992,
1088
+ "step": 1540
1089
+ },
1090
+ {
1091
+ "epoch": 1.8884826325411335,
1092
+ "grad_norm": 0.17159809172153473,
1093
+ "learning_rate": 5e-05,
1094
+ "loss": 0.7052,
1095
+ "step": 1550
1096
+ },
1097
+ {
1098
+ "epoch": 1.9006703229737965,
1099
+ "grad_norm": 0.23522134125232697,
1100
+ "learning_rate": 5e-05,
1101
+ "loss": 0.6873,
1102
+ "step": 1560
1103
+ },
1104
+ {
1105
+ "epoch": 1.9128580134064594,
1106
+ "grad_norm": 0.21959669888019562,
1107
+ "learning_rate": 5e-05,
1108
+ "loss": 0.6834,
1109
+ "step": 1570
1110
+ },
1111
+ {
1112
+ "epoch": 1.9250457038391224,
1113
+ "grad_norm": 0.2110578864812851,
1114
+ "learning_rate": 5e-05,
1115
+ "loss": 0.7038,
1116
+ "step": 1580
1117
+ },
1118
+ {
1119
+ "epoch": 1.9372333942717854,
1120
+ "grad_norm": 0.18674372136592865,
1121
+ "learning_rate": 5e-05,
1122
+ "loss": 0.6743,
1123
+ "step": 1590
1124
+ },
1125
+ {
1126
+ "epoch": 1.9494210847044484,
1127
+ "grad_norm": 0.18220046162605286,
1128
+ "learning_rate": 5e-05,
1129
+ "loss": 0.6943,
1130
+ "step": 1600
1131
+ },
1132
+ {
1133
+ "epoch": 1.9616087751371114,
1134
+ "grad_norm": 0.20756569504737854,
1135
+ "learning_rate": 5e-05,
1136
+ "loss": 0.7049,
1137
+ "step": 1610
1138
+ },
1139
+ {
1140
+ "epoch": 1.9737964655697744,
1141
+ "grad_norm": 0.1925211250782013,
1142
+ "learning_rate": 5e-05,
1143
+ "loss": 0.7033,
1144
+ "step": 1620
1145
+ },
1146
+ {
1147
+ "epoch": 1.9859841560024374,
1148
+ "grad_norm": 0.21374447643756866,
1149
+ "learning_rate": 5e-05,
1150
+ "loss": 0.7025,
1151
+ "step": 1630
1152
+ },
1153
+ {
1154
+ "epoch": 1.9981718464351006,
1155
+ "grad_norm": 0.1952880322933197,
1156
+ "learning_rate": 5e-05,
1157
+ "loss": 0.6877,
1158
+ "step": 1640
1159
+ }
1160
+ ],
1161
+ "logging_steps": 10,
1162
+ "max_steps": 2463,
1163
+ "num_input_tokens_seen": 0,
1164
+ "num_train_epochs": 3,
1165
+ "save_steps": 205,
1166
+ "stateful_callbacks": {
1167
+ "TrainerControl": {
1168
+ "args": {
1169
+ "should_epoch_stop": false,
1170
+ "should_evaluate": false,
1171
+ "should_log": false,
1172
+ "should_save": true,
1173
+ "should_training_stop": false
1174
+ },
1175
+ "attributes": {}
1176
+ }
1177
+ },
1178
+ "total_flos": 6.729135877101978e+18,
1179
+ "train_batch_size": 4,
1180
+ "trial_name": null,
1181
+ "trial_params": null
1182
+ }
openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1640/training_args.bin ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:56896ea21940e33d7e2934267dff2a5b4a4a80195a0f165245f6617a0f88b90b
3
+ size 7377
openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1640/zero_to_fp32.py ADDED
@@ -0,0 +1,760 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python
2
+
3
+ # Copyright (c) Microsoft Corporation.
4
+ # SPDX-License-Identifier: Apache-2.0
5
+
6
+ # DeepSpeed Team
7
+
8
+ # This script extracts fp32 consolidated weights from a zero 1, 2 and 3 DeepSpeed checkpoints. It gets
9
+ # copied into the top level checkpoint dir, so the user can easily do the conversion at any point in
10
+ # the future. Once extracted, the weights don't require DeepSpeed and can be used in any
11
+ # application.
12
+ #
13
+ # example:
14
+ # python zero_to_fp32.py . output_dir/
15
+ # or
16
+ # python zero_to_fp32.py . output_dir/ --safe_serialization
17
+
18
+ import argparse
19
+ import torch
20
+ import glob
21
+ import math
22
+ import os
23
+ import re
24
+ import gc
25
+ import json
26
+ import numpy as np
27
+ from tqdm import tqdm
28
+ from collections import OrderedDict
29
+ from dataclasses import dataclass
30
+
31
+ # while this script doesn't use deepspeed to recover data, since the checkpoints are pickled with
32
+ # DeepSpeed data structures it has to be available in the current python environment.
33
+ from deepspeed.utils import logger
34
+ from deepspeed.checkpoint.constants import (DS_VERSION, OPTIMIZER_STATE_DICT, SINGLE_PARTITION_OF_FP32_GROUPS,
35
+ FP32_FLAT_GROUPS, ZERO_STAGE, PARTITION_COUNT, PARAM_SHAPES, BUFFER_NAMES,
36
+ FROZEN_PARAM_SHAPES, FROZEN_PARAM_FRAGMENTS)
37
+
38
+
39
+ @dataclass
40
+ class zero_model_state:
41
+ buffers: dict()
42
+ param_shapes: dict()
43
+ shared_params: list
44
+ ds_version: int
45
+ frozen_param_shapes: dict()
46
+ frozen_param_fragments: dict()
47
+
48
+
49
+ debug = 0
50
+
51
+ # load to cpu
52
+ device = torch.device('cpu')
53
+
54
+
55
+ def atoi(text):
56
+ return int(text) if text.isdigit() else text
57
+
58
+
59
+ def natural_keys(text):
60
+ '''
61
+ alist.sort(key=natural_keys) sorts in human order
62
+ http://nedbatchelder.com/blog/200712/human_sorting.html
63
+ (See Toothy's implementation in the comments)
64
+ '''
65
+ return [atoi(c) for c in re.split(r'(\d+)', text)]
66
+
67
+
68
+ def get_model_state_file(checkpoint_dir, zero_stage):
69
+ if not os.path.isdir(checkpoint_dir):
70
+ raise FileNotFoundError(f"Directory '{checkpoint_dir}' doesn't exist")
71
+
72
+ # there should be only one file
73
+ if zero_stage <= 2:
74
+ file = os.path.join(checkpoint_dir, "mp_rank_00_model_states.pt")
75
+ elif zero_stage == 3:
76
+ file = os.path.join(checkpoint_dir, "zero_pp_rank_0_mp_rank_00_model_states.pt")
77
+
78
+ if not os.path.exists(file):
79
+ raise FileNotFoundError(f"can't find model states file at '{file}'")
80
+
81
+ return file
82
+
83
+
84
+ def get_checkpoint_files(checkpoint_dir, glob_pattern):
85
+ # XXX: need to test that this simple glob rule works for multi-node setup too
86
+ ckpt_files = sorted(glob.glob(os.path.join(checkpoint_dir, glob_pattern)), key=natural_keys)
87
+
88
+ if len(ckpt_files) == 0:
89
+ raise FileNotFoundError(f"can't find {glob_pattern} files in directory '{checkpoint_dir}'")
90
+
91
+ return ckpt_files
92
+
93
+
94
+ def get_optim_files(checkpoint_dir):
95
+ return get_checkpoint_files(checkpoint_dir, "*_optim_states.pt")
96
+
97
+
98
+ def get_model_state_files(checkpoint_dir):
99
+ return get_checkpoint_files(checkpoint_dir, "*_model_states.pt")
100
+
101
+
102
+ def parse_model_states(files):
103
+ zero_model_states = []
104
+ for file in files:
105
+ state_dict = torch.load(file, map_location=device, weights_only=False)
106
+
107
+ if BUFFER_NAMES not in state_dict:
108
+ raise ValueError(f"{file} is not a model state checkpoint")
109
+ buffer_names = state_dict[BUFFER_NAMES]
110
+ if debug:
111
+ print("Found buffers:", buffer_names)
112
+
113
+ # recover just the buffers while restoring them to fp32 if they were saved in fp16
114
+ buffers = {k: v.float() for k, v in state_dict["module"].items() if k in buffer_names}
115
+ param_shapes = state_dict[PARAM_SHAPES]
116
+
117
+ # collect parameters that are included in param_shapes
118
+ param_names = []
119
+ for s in param_shapes:
120
+ for name in s.keys():
121
+ param_names.append(name)
122
+
123
+ # update with frozen parameters
124
+ frozen_param_shapes = state_dict.get(FROZEN_PARAM_SHAPES, None)
125
+ if frozen_param_shapes is not None:
126
+ if debug:
127
+ print(f"Found frozen_param_shapes: {frozen_param_shapes}")
128
+ param_names += list(frozen_param_shapes.keys())
129
+
130
+ # handle shared params
131
+ shared_params = [[k, v] for k, v in state_dict["shared_params"].items()]
132
+
133
+ ds_version = state_dict.get(DS_VERSION, None)
134
+
135
+ frozen_param_fragments = state_dict.get(FROZEN_PARAM_FRAGMENTS, None)
136
+
137
+ z_model_state = zero_model_state(buffers=buffers,
138
+ param_shapes=param_shapes,
139
+ shared_params=shared_params,
140
+ ds_version=ds_version,
141
+ frozen_param_shapes=frozen_param_shapes,
142
+ frozen_param_fragments=frozen_param_fragments)
143
+ zero_model_states.append(z_model_state)
144
+
145
+ return zero_model_states
146
+
147
+
148
+ def parse_optim_states(files, ds_checkpoint_dir):
149
+ total_files = len(files)
150
+ state_dicts = []
151
+ for f in tqdm(files, desc='Loading checkpoint shards'):
152
+ state_dict = torch.load(f, map_location=device, mmap=True, weights_only=False)
153
+ # immediately discard the potentially huge 2 optimizer states as we only care for fp32 master weights
154
+ # and also handle the case where it was already removed by another helper script
155
+ state_dict["optimizer_state_dict"].pop("optimizer_state_dict", None)
156
+ state_dicts.append(state_dict)
157
+
158
+ if ZERO_STAGE not in state_dicts[0][OPTIMIZER_STATE_DICT]:
159
+ raise ValueError(f"{files[0]} is not a zero checkpoint")
160
+ zero_stage = state_dicts[0][OPTIMIZER_STATE_DICT][ZERO_STAGE]
161
+ world_size = state_dicts[0][OPTIMIZER_STATE_DICT][PARTITION_COUNT]
162
+
163
+ # For ZeRO-2 each param group can have different partition_count as data parallelism for expert
164
+ # parameters can be different from data parallelism for non-expert parameters. So we can just
165
+ # use the max of the partition_count to get the dp world_size.
166
+
167
+ if type(world_size) is list:
168
+ world_size = max(world_size)
169
+
170
+ if world_size != total_files:
171
+ raise ValueError(
172
+ f"Expected {world_size} of '*_optim_states.pt' under '{ds_checkpoint_dir}' but found {total_files} files. "
173
+ "Possibly due to an overwrite of an old checkpoint, or a checkpoint didn't get saved by one or more processes."
174
+ )
175
+
176
+ # the groups are named differently in each stage
177
+ if zero_stage <= 2:
178
+ fp32_groups_key = SINGLE_PARTITION_OF_FP32_GROUPS
179
+ elif zero_stage == 3:
180
+ fp32_groups_key = FP32_FLAT_GROUPS
181
+ else:
182
+ raise ValueError(f"unknown zero stage {zero_stage}")
183
+
184
+ fp32_flat_groups = [state_dicts[i][OPTIMIZER_STATE_DICT][fp32_groups_key] for i in range(len(state_dicts))]
185
+ return zero_stage, world_size, fp32_flat_groups
186
+
187
+
188
+ def _get_fp32_state_dict_from_zero_checkpoint(ds_checkpoint_dir, exclude_frozen_parameters):
189
+ """
190
+ Returns fp32 state_dict reconstructed from ds checkpoint
191
+
192
+ Args:
193
+ - ``ds_checkpoint_dir``: path to the deepspeed checkpoint folder (where the optimizer files are)
194
+
195
+ """
196
+ print(f"Processing zero checkpoint '{ds_checkpoint_dir}'")
197
+
198
+ optim_files = get_optim_files(ds_checkpoint_dir)
199
+ zero_stage, world_size, fp32_flat_groups = parse_optim_states(optim_files, ds_checkpoint_dir)
200
+ print(f"Detected checkpoint of type zero stage {zero_stage}, world_size: {world_size}")
201
+
202
+ model_files = get_model_state_files(ds_checkpoint_dir)
203
+
204
+ zero_model_states = parse_model_states(model_files)
205
+ print(f'Parsing checkpoint created by deepspeed=={zero_model_states[0].ds_version}')
206
+
207
+ if zero_stage <= 2:
208
+ return _get_fp32_state_dict_from_zero2_checkpoint(world_size, fp32_flat_groups, zero_model_states,
209
+ exclude_frozen_parameters)
210
+ elif zero_stage == 3:
211
+ return _get_fp32_state_dict_from_zero3_checkpoint(world_size, fp32_flat_groups, zero_model_states,
212
+ exclude_frozen_parameters)
213
+
214
+
215
+ def _zero2_merge_frozen_params(state_dict, zero_model_states):
216
+ if zero_model_states[0].frozen_param_shapes is None or len(zero_model_states[0].frozen_param_shapes) == 0:
217
+ return
218
+
219
+ frozen_param_shapes = zero_model_states[0].frozen_param_shapes
220
+ frozen_param_fragments = zero_model_states[0].frozen_param_fragments
221
+
222
+ if debug:
223
+ num_elem = sum(s.numel() for s in frozen_param_shapes.values())
224
+ print(f'rank 0: {FROZEN_PARAM_SHAPES}.numel = {num_elem}')
225
+
226
+ wanted_params = len(frozen_param_shapes)
227
+ wanted_numel = sum(s.numel() for s in frozen_param_shapes.values())
228
+ avail_numel = sum([p.numel() for p in frozen_param_fragments.values()])
229
+ print(f'Frozen params: Have {avail_numel} numels to process.')
230
+ print(f'Frozen params: Need {wanted_numel} numels in {wanted_params} params')
231
+
232
+ total_params = 0
233
+ total_numel = 0
234
+ for name, shape in frozen_param_shapes.items():
235
+ total_params += 1
236
+ unpartitioned_numel = shape.numel()
237
+ total_numel += unpartitioned_numel
238
+
239
+ state_dict[name] = frozen_param_fragments[name]
240
+
241
+ if debug:
242
+ print(f"{name} full shape: {shape} unpartitioned numel {unpartitioned_numel} ")
243
+
244
+ print(f"Reconstructed Frozen fp32 state dict with {total_params} params {total_numel} elements")
245
+
246
+
247
+ def _has_callable(obj, fn):
248
+ attr = getattr(obj, fn, None)
249
+ return callable(attr)
250
+
251
+
252
+ def _zero2_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states):
253
+ param_shapes = zero_model_states[0].param_shapes
254
+
255
+ # Reconstruction protocol:
256
+ #
257
+ # XXX: document this
258
+
259
+ if debug:
260
+ for i in range(world_size):
261
+ for j in range(len(fp32_flat_groups[0])):
262
+ print(f"{FP32_FLAT_GROUPS}[{i}][{j}].shape={fp32_flat_groups[i][j].shape}")
263
+
264
+ # XXX: memory usage doubles here (zero2)
265
+ num_param_groups = len(fp32_flat_groups[0])
266
+ merged_single_partition_of_fp32_groups = []
267
+ for i in range(num_param_groups):
268
+ merged_partitions = [sd[i] for sd in fp32_flat_groups]
269
+ full_single_fp32_vector = torch.cat(merged_partitions, 0)
270
+ merged_single_partition_of_fp32_groups.append(full_single_fp32_vector)
271
+ avail_numel = sum(
272
+ [full_single_fp32_vector.numel() for full_single_fp32_vector in merged_single_partition_of_fp32_groups])
273
+
274
+ if debug:
275
+ wanted_params = sum([len(shapes) for shapes in param_shapes])
276
+ wanted_numel = sum([sum(shape.numel() for shape in shapes.values()) for shapes in param_shapes])
277
+ # not asserting if there is a mismatch due to possible padding
278
+ print(f"Have {avail_numel} numels to process.")
279
+ print(f"Need {wanted_numel} numels in {wanted_params} params.")
280
+
281
+ # params
282
+ # XXX: for huge models that can't fit into the host's RAM we will have to recode this to support
283
+ # out-of-core computing solution
284
+ total_numel = 0
285
+ total_params = 0
286
+ for shapes, full_single_fp32_vector in zip(param_shapes, merged_single_partition_of_fp32_groups):
287
+ offset = 0
288
+ avail_numel = full_single_fp32_vector.numel()
289
+ for name, shape in shapes.items():
290
+
291
+ unpartitioned_numel = shape.numel() if _has_callable(shape, 'numel') else math.prod(shape)
292
+ total_numel += unpartitioned_numel
293
+ total_params += 1
294
+
295
+ if debug:
296
+ print(f"{name} full shape: {shape} unpartitioned numel {unpartitioned_numel} ")
297
+ state_dict[name] = full_single_fp32_vector.narrow(0, offset, unpartitioned_numel).view(shape)
298
+ offset += unpartitioned_numel
299
+
300
+ # Z2 started to align to 2*world_size to improve nccl performance. Therefore both offset and
301
+ # avail_numel can differ by anywhere between 0..2*world_size. Due to two unrelated complex
302
+ # paddings performed in the code it's almost impossible to predict the exact numbers w/o the
303
+ # live optimizer object, so we are checking that the numbers are within the right range
304
+ align_to = 2 * world_size
305
+
306
+ def zero2_align(x):
307
+ return align_to * math.ceil(x / align_to)
308
+
309
+ if debug:
310
+ print(f"original offset={offset}, avail_numel={avail_numel}")
311
+
312
+ offset = zero2_align(offset)
313
+ avail_numel = zero2_align(avail_numel)
314
+
315
+ if debug:
316
+ print(f"aligned offset={offset}, avail_numel={avail_numel}")
317
+
318
+ # Sanity check
319
+ if offset != avail_numel:
320
+ raise ValueError(f"consumed {offset} numels out of {avail_numel} - something is wrong")
321
+
322
+ print(f"Reconstructed fp32 state dict with {total_params} params {total_numel} elements")
323
+
324
+
325
+ def _get_fp32_state_dict_from_zero2_checkpoint(world_size, fp32_flat_groups, zero_model_states,
326
+ exclude_frozen_parameters):
327
+ state_dict = OrderedDict()
328
+
329
+ # buffers
330
+ buffers = zero_model_states[0].buffers
331
+ state_dict.update(buffers)
332
+ if debug:
333
+ print(f"added {len(buffers)} buffers")
334
+
335
+ if not exclude_frozen_parameters:
336
+ _zero2_merge_frozen_params(state_dict, zero_model_states)
337
+
338
+ _zero2_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states)
339
+
340
+ # recover shared parameters
341
+ for pair in zero_model_states[0].shared_params:
342
+ if pair[1] in state_dict:
343
+ state_dict[pair[0]] = state_dict[pair[1]]
344
+
345
+ return state_dict
346
+
347
+
348
+ def zero3_partitioned_param_info(unpartitioned_numel, world_size):
349
+ remainder = unpartitioned_numel % world_size
350
+ padding_numel = (world_size - remainder) if remainder else 0
351
+ partitioned_numel = math.ceil(unpartitioned_numel / world_size)
352
+ return partitioned_numel, padding_numel
353
+
354
+
355
+ def _zero3_merge_frozen_params(state_dict, world_size, zero_model_states):
356
+ if zero_model_states[0].frozen_param_shapes is None or len(zero_model_states[0].frozen_param_shapes) == 0:
357
+ return
358
+
359
+ if debug:
360
+ for i in range(world_size):
361
+ num_elem = sum(s.numel() for s in zero_model_states[i].frozen_param_fragments.values())
362
+ print(f'rank {i}: {FROZEN_PARAM_SHAPES}.numel = {num_elem}')
363
+
364
+ frozen_param_shapes = zero_model_states[0].frozen_param_shapes
365
+ wanted_params = len(frozen_param_shapes)
366
+ wanted_numel = sum(s.numel() for s in frozen_param_shapes.values())
367
+ avail_numel = sum([p.numel() for p in zero_model_states[0].frozen_param_fragments.values()]) * world_size
368
+ print(f'Frozen params: Have {avail_numel} numels to process.')
369
+ print(f'Frozen params: Need {wanted_numel} numels in {wanted_params} params')
370
+
371
+ total_params = 0
372
+ total_numel = 0
373
+ for name, shape in zero_model_states[0].frozen_param_shapes.items():
374
+ total_params += 1
375
+ unpartitioned_numel = shape.numel()
376
+ total_numel += unpartitioned_numel
377
+
378
+ param_frags = tuple(model_state.frozen_param_fragments[name] for model_state in zero_model_states)
379
+ state_dict[name] = torch.cat(param_frags, 0).narrow(0, 0, unpartitioned_numel).view(shape)
380
+
381
+ partitioned_numel, partitioned_padding_numel = zero3_partitioned_param_info(unpartitioned_numel, world_size)
382
+
383
+ if debug:
384
+ print(
385
+ f"Frozen params: {total_params} {name} full shape: {shape} partition0 numel={partitioned_numel} partitioned_padding_numel={partitioned_padding_numel}"
386
+ )
387
+
388
+ print(f"Reconstructed Frozen fp32 state dict with {total_params} params {total_numel} elements")
389
+
390
+
391
+ class GatheredTensor:
392
+ """
393
+ A pseudo tensor that collects partitioned weights.
394
+ It is more memory efficient when there are multiple groups.
395
+ """
396
+
397
+ def __init__(self, flat_groups, flat_groups_offset, offset, partitioned_numel, shape):
398
+ self.flat_groups = flat_groups
399
+ self.flat_groups_offset = flat_groups_offset
400
+ self.offset = offset
401
+ self.partitioned_numel = partitioned_numel
402
+ self.shape = shape
403
+ self.dtype = self.flat_groups[0][0].dtype
404
+
405
+ def contiguous(self):
406
+ """
407
+ Merge partitioned weights from flat_groups into a single tensor.
408
+ """
409
+ end_idx = self.offset + self.partitioned_numel
410
+ world_size = len(self.flat_groups)
411
+ pad_flat_param_chunks = []
412
+
413
+ for rank_i in range(world_size):
414
+ # for each rank, we need to collect weights from related group/groups
415
+ flat_groups_at_rank_i = self.flat_groups[rank_i]
416
+ start_group_id = None
417
+ end_group_id = None
418
+ for group_id in range(len(self.flat_groups_offset)):
419
+ if self.flat_groups_offset[group_id] <= self.offset < self.flat_groups_offset[group_id + 1]:
420
+ start_group_id = group_id
421
+ if self.flat_groups_offset[group_id] < end_idx <= self.flat_groups_offset[group_id + 1]:
422
+ end_group_id = group_id
423
+ break
424
+ # collect weights from related group/groups
425
+ for group_id in range(start_group_id, end_group_id + 1):
426
+ flat_tensor = flat_groups_at_rank_i[group_id]
427
+ start_offset = self.offset - self.flat_groups_offset[group_id]
428
+ end_offset = min(end_idx, self.flat_groups_offset[group_id + 1]) - self.flat_groups_offset[group_id]
429
+ pad_flat_param_chunks.append(flat_tensor[start_offset:end_offset])
430
+
431
+ # collect weights from all ranks
432
+ pad_flat_param = torch.cat(pad_flat_param_chunks, dim=0)
433
+ param = pad_flat_param[:self.shape.numel()].view(self.shape).contiguous()
434
+ return param
435
+
436
+
437
+ def _zero3_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states):
438
+ param_shapes = zero_model_states[0].param_shapes
439
+ avail_numel = sum([flat_group.numel() for flat_group in fp32_flat_groups[0]]) * world_size
440
+
441
+ # Reconstruction protocol: For zero3 we need to zip the partitions together at boundary of each
442
+ # param, re-consolidating each param, while dealing with padding if any
443
+
444
+ # merge list of dicts, preserving order
445
+ param_shapes = {k: v for d in param_shapes for k, v in d.items()}
446
+
447
+ if debug:
448
+ for i in range(world_size):
449
+ print(f"{FP32_FLAT_GROUPS}[{i}].shape={fp32_flat_groups[i].shape}")
450
+
451
+ wanted_params = len(param_shapes)
452
+ wanted_numel = sum(shape.numel() for shape in param_shapes.values())
453
+ # not asserting if there is a mismatch due to possible padding
454
+ avail_numel = fp32_flat_groups[0].numel() * world_size
455
+ print(f"Trainable params: Have {avail_numel} numels to process.")
456
+ print(f"Trainable params: Need {wanted_numel} numels in {wanted_params} params.")
457
+
458
+ # params
459
+ # XXX: for huge models that can't fit into the host's RAM we will have to recode this to support
460
+ # out-of-core computing solution
461
+ offset = 0
462
+ total_numel = 0
463
+ total_params = 0
464
+ flat_groups_offset = [0] + list(np.cumsum([flat_tensor.numel() for flat_tensor in fp32_flat_groups[0]]))
465
+ for name, shape in tqdm(param_shapes.items(), desc='Gathering sharded weights'):
466
+ unpartitioned_numel = shape.numel()
467
+ total_numel += unpartitioned_numel
468
+ total_params += 1
469
+ partitioned_numel, partitioned_padding_numel = zero3_partitioned_param_info(unpartitioned_numel, world_size)
470
+
471
+ if debug:
472
+ print(
473
+ f"Trainable params: {total_params} {name} full shape: {shape} partition0 numel={partitioned_numel} partitioned_padding_numel={partitioned_padding_numel}"
474
+ )
475
+
476
+ # memory efficient tensor
477
+ tensor = GatheredTensor(fp32_flat_groups, flat_groups_offset, offset, partitioned_numel, shape)
478
+ state_dict[name] = tensor
479
+ offset += partitioned_numel
480
+
481
+ offset *= world_size
482
+
483
+ # Sanity check
484
+ if offset != avail_numel:
485
+ raise ValueError(f"consumed {offset} numels out of {avail_numel} - something is wrong")
486
+
487
+ print(f"Reconstructed Trainable fp32 state dict with {total_params} params {total_numel} elements")
488
+
489
+
490
+ def _get_fp32_state_dict_from_zero3_checkpoint(world_size, fp32_flat_groups, zero_model_states,
491
+ exclude_frozen_parameters):
492
+ state_dict = OrderedDict()
493
+
494
+ # buffers
495
+ buffers = zero_model_states[0].buffers
496
+ state_dict.update(buffers)
497
+ if debug:
498
+ print(f"added {len(buffers)} buffers")
499
+
500
+ if not exclude_frozen_parameters:
501
+ _zero3_merge_frozen_params(state_dict, world_size, zero_model_states)
502
+
503
+ _zero3_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states)
504
+
505
+ # recover shared parameters
506
+ for pair in zero_model_states[0].shared_params:
507
+ if pair[1] in state_dict:
508
+ state_dict[pair[0]] = state_dict[pair[1]]
509
+
510
+ return state_dict
511
+
512
+
513
+ def to_torch_tensor(state_dict, return_empty_tensor=False):
514
+ """
515
+ Convert state_dict of GatheredTensor to torch tensor
516
+ """
517
+ torch_state_dict = {}
518
+ converted_tensors = {}
519
+ for name, tensor in state_dict.items():
520
+ tensor_id = id(tensor)
521
+ if tensor_id in converted_tensors: # shared tensors
522
+ shared_tensor = torch_state_dict[converted_tensors[tensor_id]]
523
+ torch_state_dict[name] = shared_tensor
524
+ else:
525
+ converted_tensors[tensor_id] = name
526
+ if return_empty_tensor:
527
+ torch_state_dict[name] = torch.empty(tensor.shape, dtype=tensor.dtype)
528
+ else:
529
+ torch_state_dict[name] = tensor.contiguous()
530
+ return torch_state_dict
531
+
532
+
533
+ def get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir,
534
+ tag=None,
535
+ exclude_frozen_parameters=False,
536
+ lazy_mode=False):
537
+ """
538
+ Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated state_dict that can be loaded with
539
+ ``load_state_dict()`` and used for training without DeepSpeed or shared with others, for example
540
+ via a model hub.
541
+
542
+ Args:
543
+ - ``checkpoint_dir``: path to the desired checkpoint folder
544
+ - ``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``
545
+ - ``exclude_frozen_parameters``: exclude frozen parameters
546
+ - ``lazy_mode``: get state_dict in lazy mode. It returns a dict of pesduo tensor instead of torch tensor, which is more memory efficient.
547
+ Convert the pesduo tensor to torch tensor by ``.contiguous()``
548
+
549
+ Returns:
550
+ - pytorch ``state_dict``
551
+
552
+ A typical usage might be ::
553
+
554
+ from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint
555
+ # do the training and checkpoint saving
556
+ state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir) # already on cpu
557
+ model = model.cpu() # move to cpu
558
+ model.load_state_dict(state_dict)
559
+ # submit to model hub or save the model to share with others
560
+
561
+ In this example the ``model`` will no longer be usable in the deepspeed context of the same
562
+ application. i.e. you will need to re-initialize the deepspeed engine, since
563
+ ``model.load_state_dict(state_dict)`` will remove all the deepspeed magic from it.
564
+
565
+ If you want it all done for you, use ``load_state_dict_from_zero_checkpoint`` instead.
566
+
567
+ Note: the above usage may not work if your application doesn't have sufficient free CPU memory.
568
+ You may need to use the offline approach using the ``zero_to_fp32.py`` script that is saved with
569
+ the checkpoint. Or you can load state_dict in lazy mode ::
570
+
571
+ from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint
572
+ state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, lazy_mode=True) # not on cpu
573
+ for name, lazy_tensor in state_dict.item():
574
+ tensor = lazy_tensor.contiguous() # to cpu
575
+ print(name, tensor)
576
+ # del tensor to release memory if it no longer in use
577
+ """
578
+ if tag is None:
579
+ latest_path = os.path.join(checkpoint_dir, 'latest')
580
+ if os.path.isfile(latest_path):
581
+ with open(latest_path, 'r') as fd:
582
+ tag = fd.read().strip()
583
+ else:
584
+ raise ValueError(f"Unable to find 'latest' file at {latest_path}")
585
+
586
+ ds_checkpoint_dir = os.path.join(checkpoint_dir, tag)
587
+
588
+ if not os.path.isdir(ds_checkpoint_dir):
589
+ raise FileNotFoundError(f"Directory '{ds_checkpoint_dir}' doesn't exist")
590
+
591
+ state_dict = _get_fp32_state_dict_from_zero_checkpoint(ds_checkpoint_dir, exclude_frozen_parameters)
592
+ if lazy_mode:
593
+ return state_dict
594
+ else:
595
+ return to_torch_tensor(state_dict)
596
+
597
+
598
+ def convert_zero_checkpoint_to_fp32_state_dict(checkpoint_dir,
599
+ output_dir,
600
+ max_shard_size="5GB",
601
+ safe_serialization=False,
602
+ tag=None,
603
+ exclude_frozen_parameters=False):
604
+ """
605
+ Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated ``state_dict`` file that can be
606
+ loaded with ``torch.load(file)`` + ``load_state_dict()`` and used for training without DeepSpeed.
607
+
608
+ Args:
609
+ - ``checkpoint_dir``: path to the desired checkpoint folder. (one that contains the tag-folder, like ``global_step14``)
610
+ - ``output_dir``: directory to the pytorch fp32 state_dict output files
611
+ - ``max_shard_size``: the maximum size for a checkpoint before being sharded, default value is 5GB
612
+ - ``safe_serialization``: whether to save the model using `safetensors` or the traditional PyTorch way (that uses `pickle`).
613
+ - ``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``
614
+ - ``exclude_frozen_parameters``: exclude frozen parameters
615
+ """
616
+
617
+ # Dependency pre-check
618
+ if safe_serialization:
619
+ try:
620
+ from safetensors.torch import save_file
621
+ except ImportError:
622
+ print('If you want to use `safe_serialization`, please `pip install safetensors`')
623
+ raise
624
+ if max_shard_size is not None:
625
+ try:
626
+ from huggingface_hub import split_torch_state_dict_into_shards
627
+ except ImportError:
628
+ print('If you want to use `max_shard_size`, please `pip install huggingface_hub`')
629
+ raise
630
+
631
+ # Convert zero checkpoint to state_dict
632
+ state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir,
633
+ tag,
634
+ exclude_frozen_parameters,
635
+ lazy_mode=True)
636
+
637
+ # Shard the model if it is too big.
638
+ weights_name = "model.safetensors" if safe_serialization else "pytorch_model.bin"
639
+ if max_shard_size is not None:
640
+ filename_pattern = weights_name.replace(".bin", "{suffix}.bin").replace(".safetensors", "{suffix}.safetensors")
641
+ # an memory-efficient approach for sharding
642
+ empty_state_dict = to_torch_tensor(state_dict, return_empty_tensor=True)
643
+ state_dict_split = split_torch_state_dict_into_shards(empty_state_dict,
644
+ filename_pattern=filename_pattern,
645
+ max_shard_size=max_shard_size)
646
+ else:
647
+ from collections import namedtuple
648
+ StateDictSplit = namedtuple("StateDictSplit", ["is_sharded", "filename_to_tensors"])
649
+ state_dict_split = StateDictSplit(is_sharded=False,
650
+ filename_to_tensors={weights_name: list(state_dict.keys())})
651
+
652
+ # Save the model by shard
653
+ os.makedirs(output_dir, exist_ok=True)
654
+ filename_to_tensors = state_dict_split.filename_to_tensors.items()
655
+ for shard_file, tensors in tqdm(filename_to_tensors, desc="Saving checkpoint shards"):
656
+ shard_state_dict = {tensor_name: state_dict[tensor_name] for tensor_name in tensors}
657
+ shard_state_dict = to_torch_tensor(shard_state_dict)
658
+ output_path = os.path.join(output_dir, shard_file)
659
+ if safe_serialization:
660
+ save_file(shard_state_dict, output_path, metadata={"format": "pt"})
661
+ else:
662
+ torch.save(shard_state_dict, output_path)
663
+ # release the memory of current shard
664
+ for tensor_name in list(shard_state_dict.keys()):
665
+ del state_dict[tensor_name]
666
+ del shard_state_dict[tensor_name]
667
+ del shard_state_dict
668
+ gc.collect()
669
+
670
+ # Save index if sharded
671
+ if state_dict_split.is_sharded:
672
+ index = {
673
+ "metadata": state_dict_split.metadata,
674
+ "weight_map": state_dict_split.tensor_to_filename,
675
+ }
676
+ save_index_file = "model.safetensors.index.json" if safe_serialization else "pytorch_model.bin.index.json"
677
+ save_index_file = os.path.join(output_dir, save_index_file)
678
+ with open(save_index_file, "w", encoding="utf-8") as f:
679
+ content = json.dumps(index, indent=2, sort_keys=True) + "\n"
680
+ f.write(content)
681
+
682
+
683
+ def load_state_dict_from_zero_checkpoint(model, checkpoint_dir, tag=None):
684
+ """
685
+ 1. Put the provided model to cpu
686
+ 2. Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated ``state_dict``
687
+ 3. Load it into the provided model
688
+
689
+ Args:
690
+ - ``model``: the model object to update
691
+ - ``checkpoint_dir``: path to the desired checkpoint folder. (one that contains the tag-folder, like ``global_step14``)
692
+ - ``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``
693
+
694
+ Returns:
695
+ - ``model`: modified model
696
+
697
+ Make sure you have plenty of CPU memory available before you call this function. If you don't
698
+ have enough use the ``zero_to_fp32.py`` utility to do the conversion. You will find it
699
+ conveniently placed for you in the checkpoint folder.
700
+
701
+ A typical usage might be ::
702
+
703
+ from deepspeed.utils.zero_to_fp32 import load_state_dict_from_zero_checkpoint
704
+ model = load_state_dict_from_zero_checkpoint(trainer.model, checkpoint_dir)
705
+ # submit to model hub or save the model to share with others
706
+
707
+ Note, that once this was run, the ``model`` will no longer be usable in the deepspeed context
708
+ of the same application. i.e. you will need to re-initialize the deepspeed engine, since
709
+ ``model.load_state_dict(state_dict)`` will remove all the deepspeed magic from it.
710
+
711
+ """
712
+ logger.info("Extracting fp32 weights")
713
+ state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, tag)
714
+
715
+ logger.info("Overwriting model with fp32 weights")
716
+ model = model.cpu()
717
+ model.load_state_dict(state_dict, strict=False)
718
+
719
+ return model
720
+
721
+
722
+ if __name__ == "__main__":
723
+ parser = argparse.ArgumentParser()
724
+ parser.add_argument("checkpoint_dir",
725
+ type=str,
726
+ help="path to the desired checkpoint folder, e.g., path/checkpoint-12")
727
+ parser.add_argument("output_dir",
728
+ type=str,
729
+ help="directory to the pytorch fp32 state_dict output files"
730
+ "(e.g. path/checkpoint-12-output/)")
731
+ parser.add_argument(
732
+ "--max_shard_size",
733
+ type=str,
734
+ default="5GB",
735
+ help="The maximum size for a checkpoint before being sharded. Checkpoints shard will then be each of size"
736
+ "lower than this size. If expressed as a string, needs to be digits followed by a unit (like `5MB`"
737
+ "We default it to 5GB in order for models to be able to run easily on free-tier google colab instances"
738
+ "without CPU OOM issues.")
739
+ parser.add_argument(
740
+ "--safe_serialization",
741
+ default=False,
742
+ action='store_true',
743
+ help="Whether to save the model using `safetensors` or the traditional PyTorch way (that uses `pickle`).")
744
+ parser.add_argument("-t",
745
+ "--tag",
746
+ type=str,
747
+ default=None,
748
+ help="checkpoint tag used as a unique identifier for checkpoint. e.g., global_step1")
749
+ parser.add_argument("--exclude_frozen_parameters", action='store_true', help="exclude frozen parameters")
750
+ parser.add_argument("-d", "--debug", action='store_true', help="enable debug")
751
+ args = parser.parse_args()
752
+
753
+ debug = args.debug
754
+
755
+ convert_zero_checkpoint_to_fp32_state_dict(args.checkpoint_dir,
756
+ args.output_dir,
757
+ max_shard_size=args.max_shard_size,
758
+ safe_serialization=args.safe_serialization,
759
+ tag=args.tag,
760
+ exclude_frozen_parameters=args.exclude_frozen_parameters)
openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1845/config.json ADDED
@@ -0,0 +1,39 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "architectures": [
3
+ "LlamaForCausalLM"
4
+ ],
5
+ "attention_bias": false,
6
+ "attention_dropout": 0.0,
7
+ "bos_token_id": 128000,
8
+ "dtype": "bfloat16",
9
+ "eos_token_id": [
10
+ 128001,
11
+ 128008,
12
+ 128009
13
+ ],
14
+ "head_dim": 64,
15
+ "hidden_act": "silu",
16
+ "hidden_size": 2048,
17
+ "initializer_range": 0.02,
18
+ "intermediate_size": 8192,
19
+ "max_position_embeddings": 131072,
20
+ "mlp_bias": false,
21
+ "model_type": "llama",
22
+ "num_attention_heads": 32,
23
+ "num_hidden_layers": 16,
24
+ "num_key_value_heads": 8,
25
+ "pretraining_tp": 1,
26
+ "rms_norm_eps": 1e-05,
27
+ "rope_scaling": {
28
+ "factor": 32.0,
29
+ "high_freq_factor": 4.0,
30
+ "low_freq_factor": 1.0,
31
+ "original_max_position_embeddings": 8192,
32
+ "rope_type": "llama3"
33
+ },
34
+ "rope_theta": 500000.0,
35
+ "tie_word_embeddings": true,
36
+ "transformers_version": "4.57.1",
37
+ "use_cache": false,
38
+ "vocab_size": 128256
39
+ }
openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1845/generation_config.json ADDED
@@ -0,0 +1,12 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "bos_token_id": 128000,
3
+ "do_sample": true,
4
+ "eos_token_id": [
5
+ 128001,
6
+ 128008,
7
+ 128009
8
+ ],
9
+ "temperature": 0.6,
10
+ "top_p": 0.9,
11
+ "transformers_version": "4.57.1"
12
+ }
openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1845/latest ADDED
@@ -0,0 +1 @@
 
 
1
+ global_step1845
openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1845/model.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:4c2c95dc323c92307bbea2ca9c79b15a8e18c558f447ae69646eaecd7ca3fec5
3
+ size 2996982344
openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1845/scheduler.pt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:45bf000fd03fbe6a60b9b8cbe4f9d2b1df7587db393fc2d83075b5fc0b286922
3
+ size 1465
openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1845/trainer_state.json ADDED
@@ -0,0 +1,1322 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "best_global_step": null,
3
+ "best_metric": null,
4
+ "best_model_checkpoint": null,
5
+ "epoch": 2.247410115783059,
6
+ "eval_steps": 500,
7
+ "global_step": 1845,
8
+ "is_hyper_param_search": false,
9
+ "is_local_process_zero": true,
10
+ "is_world_process_zero": true,
11
+ "log_history": [
12
+ {
13
+ "epoch": 0.01218769043266301,
14
+ "grad_norm": 0.36486390233039856,
15
+ "learning_rate": 5e-05,
16
+ "loss": 1.2213,
17
+ "step": 10
18
+ },
19
+ {
20
+ "epoch": 0.02437538086532602,
21
+ "grad_norm": 0.2953564524650574,
22
+ "learning_rate": 5e-05,
23
+ "loss": 1.0223,
24
+ "step": 20
25
+ },
26
+ {
27
+ "epoch": 0.03656307129798903,
28
+ "grad_norm": 0.2766735851764679,
29
+ "learning_rate": 5e-05,
30
+ "loss": 0.9589,
31
+ "step": 30
32
+ },
33
+ {
34
+ "epoch": 0.04875076173065204,
35
+ "grad_norm": 0.32770997285842896,
36
+ "learning_rate": 5e-05,
37
+ "loss": 0.9081,
38
+ "step": 40
39
+ },
40
+ {
41
+ "epoch": 0.06093845216331505,
42
+ "grad_norm": 0.3052276372909546,
43
+ "learning_rate": 5e-05,
44
+ "loss": 0.8752,
45
+ "step": 50
46
+ },
47
+ {
48
+ "epoch": 0.07312614259597806,
49
+ "grad_norm": 0.240146666765213,
50
+ "learning_rate": 5e-05,
51
+ "loss": 0.873,
52
+ "step": 60
53
+ },
54
+ {
55
+ "epoch": 0.08531383302864107,
56
+ "grad_norm": 0.23872187733650208,
57
+ "learning_rate": 5e-05,
58
+ "loss": 0.8587,
59
+ "step": 70
60
+ },
61
+ {
62
+ "epoch": 0.09750152346130408,
63
+ "grad_norm": 0.32638487219810486,
64
+ "learning_rate": 5e-05,
65
+ "loss": 0.8656,
66
+ "step": 80
67
+ },
68
+ {
69
+ "epoch": 0.10968921389396709,
70
+ "grad_norm": 0.25019028782844543,
71
+ "learning_rate": 5e-05,
72
+ "loss": 0.8685,
73
+ "step": 90
74
+ },
75
+ {
76
+ "epoch": 0.1218769043266301,
77
+ "grad_norm": 0.2507675886154175,
78
+ "learning_rate": 5e-05,
79
+ "loss": 0.8662,
80
+ "step": 100
81
+ },
82
+ {
83
+ "epoch": 0.1340645947592931,
84
+ "grad_norm": 0.27966755628585815,
85
+ "learning_rate": 5e-05,
86
+ "loss": 0.8428,
87
+ "step": 110
88
+ },
89
+ {
90
+ "epoch": 0.14625228519195613,
91
+ "grad_norm": 0.2622673213481903,
92
+ "learning_rate": 5e-05,
93
+ "loss": 0.8416,
94
+ "step": 120
95
+ },
96
+ {
97
+ "epoch": 0.15843997562461914,
98
+ "grad_norm": 0.3003966510295868,
99
+ "learning_rate": 5e-05,
100
+ "loss": 0.8528,
101
+ "step": 130
102
+ },
103
+ {
104
+ "epoch": 0.17062766605728213,
105
+ "grad_norm": 0.32308319211006165,
106
+ "learning_rate": 5e-05,
107
+ "loss": 0.8357,
108
+ "step": 140
109
+ },
110
+ {
111
+ "epoch": 0.18281535648994515,
112
+ "grad_norm": 0.27279096841812134,
113
+ "learning_rate": 5e-05,
114
+ "loss": 0.8432,
115
+ "step": 150
116
+ },
117
+ {
118
+ "epoch": 0.19500304692260817,
119
+ "grad_norm": 0.27976104617118835,
120
+ "learning_rate": 5e-05,
121
+ "loss": 0.8119,
122
+ "step": 160
123
+ },
124
+ {
125
+ "epoch": 0.2071907373552712,
126
+ "grad_norm": 0.290866881608963,
127
+ "learning_rate": 5e-05,
128
+ "loss": 0.8096,
129
+ "step": 170
130
+ },
131
+ {
132
+ "epoch": 0.21937842778793418,
133
+ "grad_norm": 0.29430949687957764,
134
+ "learning_rate": 5e-05,
135
+ "loss": 0.819,
136
+ "step": 180
137
+ },
138
+ {
139
+ "epoch": 0.2315661182205972,
140
+ "grad_norm": 0.27714136242866516,
141
+ "learning_rate": 5e-05,
142
+ "loss": 0.8187,
143
+ "step": 190
144
+ },
145
+ {
146
+ "epoch": 0.2437538086532602,
147
+ "grad_norm": 0.2997029423713684,
148
+ "learning_rate": 5e-05,
149
+ "loss": 0.808,
150
+ "step": 200
151
+ },
152
+ {
153
+ "epoch": 0.25594149908592323,
154
+ "grad_norm": 0.2539573609828949,
155
+ "learning_rate": 5e-05,
156
+ "loss": 0.7943,
157
+ "step": 210
158
+ },
159
+ {
160
+ "epoch": 0.2681291895185862,
161
+ "grad_norm": 0.23294171690940857,
162
+ "learning_rate": 5e-05,
163
+ "loss": 0.8188,
164
+ "step": 220
165
+ },
166
+ {
167
+ "epoch": 0.28031687995124926,
168
+ "grad_norm": 0.25016292929649353,
169
+ "learning_rate": 5e-05,
170
+ "loss": 0.8035,
171
+ "step": 230
172
+ },
173
+ {
174
+ "epoch": 0.29250457038391225,
175
+ "grad_norm": 0.37125322222709656,
176
+ "learning_rate": 5e-05,
177
+ "loss": 0.785,
178
+ "step": 240
179
+ },
180
+ {
181
+ "epoch": 0.30469226081657524,
182
+ "grad_norm": 0.21520452201366425,
183
+ "learning_rate": 5e-05,
184
+ "loss": 0.8088,
185
+ "step": 250
186
+ },
187
+ {
188
+ "epoch": 0.3168799512492383,
189
+ "grad_norm": 0.20990602672100067,
190
+ "learning_rate": 5e-05,
191
+ "loss": 0.7813,
192
+ "step": 260
193
+ },
194
+ {
195
+ "epoch": 0.3290676416819013,
196
+ "grad_norm": 0.24193961918354034,
197
+ "learning_rate": 5e-05,
198
+ "loss": 0.7979,
199
+ "step": 270
200
+ },
201
+ {
202
+ "epoch": 0.34125533211456427,
203
+ "grad_norm": 0.22766050696372986,
204
+ "learning_rate": 5e-05,
205
+ "loss": 0.785,
206
+ "step": 280
207
+ },
208
+ {
209
+ "epoch": 0.3534430225472273,
210
+ "grad_norm": 0.2662702202796936,
211
+ "learning_rate": 5e-05,
212
+ "loss": 0.7766,
213
+ "step": 290
214
+ },
215
+ {
216
+ "epoch": 0.3656307129798903,
217
+ "grad_norm": 0.23020495474338531,
218
+ "learning_rate": 5e-05,
219
+ "loss": 0.7687,
220
+ "step": 300
221
+ },
222
+ {
223
+ "epoch": 0.3778184034125533,
224
+ "grad_norm": 0.21770520508289337,
225
+ "learning_rate": 5e-05,
226
+ "loss": 0.7815,
227
+ "step": 310
228
+ },
229
+ {
230
+ "epoch": 0.39000609384521634,
231
+ "grad_norm": 0.22677403688430786,
232
+ "learning_rate": 5e-05,
233
+ "loss": 0.7999,
234
+ "step": 320
235
+ },
236
+ {
237
+ "epoch": 0.40219378427787933,
238
+ "grad_norm": 0.24891860783100128,
239
+ "learning_rate": 5e-05,
240
+ "loss": 0.7946,
241
+ "step": 330
242
+ },
243
+ {
244
+ "epoch": 0.4143814747105424,
245
+ "grad_norm": 0.2506854832172394,
246
+ "learning_rate": 5e-05,
247
+ "loss": 0.7889,
248
+ "step": 340
249
+ },
250
+ {
251
+ "epoch": 0.42656916514320536,
252
+ "grad_norm": 0.24388180673122406,
253
+ "learning_rate": 5e-05,
254
+ "loss": 0.79,
255
+ "step": 350
256
+ },
257
+ {
258
+ "epoch": 0.43875685557586835,
259
+ "grad_norm": 0.22329650819301605,
260
+ "learning_rate": 5e-05,
261
+ "loss": 0.7632,
262
+ "step": 360
263
+ },
264
+ {
265
+ "epoch": 0.4509445460085314,
266
+ "grad_norm": 0.2240927815437317,
267
+ "learning_rate": 5e-05,
268
+ "loss": 0.8042,
269
+ "step": 370
270
+ },
271
+ {
272
+ "epoch": 0.4631322364411944,
273
+ "grad_norm": 0.2578203082084656,
274
+ "learning_rate": 5e-05,
275
+ "loss": 0.7667,
276
+ "step": 380
277
+ },
278
+ {
279
+ "epoch": 0.4753199268738574,
280
+ "grad_norm": 0.25349217653274536,
281
+ "learning_rate": 5e-05,
282
+ "loss": 0.7911,
283
+ "step": 390
284
+ },
285
+ {
286
+ "epoch": 0.4875076173065204,
287
+ "grad_norm": 0.254446804523468,
288
+ "learning_rate": 5e-05,
289
+ "loss": 0.7909,
290
+ "step": 400
291
+ },
292
+ {
293
+ "epoch": 0.4996953077391834,
294
+ "grad_norm": 0.24253171682357788,
295
+ "learning_rate": 5e-05,
296
+ "loss": 0.7896,
297
+ "step": 410
298
+ },
299
+ {
300
+ "epoch": 0.5118829981718465,
301
+ "grad_norm": 0.27901458740234375,
302
+ "learning_rate": 5e-05,
303
+ "loss": 0.7828,
304
+ "step": 420
305
+ },
306
+ {
307
+ "epoch": 0.5240706886045094,
308
+ "grad_norm": 0.31884434819221497,
309
+ "learning_rate": 5e-05,
310
+ "loss": 0.7919,
311
+ "step": 430
312
+ },
313
+ {
314
+ "epoch": 0.5362583790371724,
315
+ "grad_norm": 0.22920799255371094,
316
+ "learning_rate": 5e-05,
317
+ "loss": 0.7777,
318
+ "step": 440
319
+ },
320
+ {
321
+ "epoch": 0.5484460694698354,
322
+ "grad_norm": 0.22606618702411652,
323
+ "learning_rate": 5e-05,
324
+ "loss": 0.7661,
325
+ "step": 450
326
+ },
327
+ {
328
+ "epoch": 0.5606337599024985,
329
+ "grad_norm": 0.21604052186012268,
330
+ "learning_rate": 5e-05,
331
+ "loss": 0.769,
332
+ "step": 460
333
+ },
334
+ {
335
+ "epoch": 0.5728214503351615,
336
+ "grad_norm": 0.2296101152896881,
337
+ "learning_rate": 5e-05,
338
+ "loss": 0.7626,
339
+ "step": 470
340
+ },
341
+ {
342
+ "epoch": 0.5850091407678245,
343
+ "grad_norm": 0.24162150919437408,
344
+ "learning_rate": 5e-05,
345
+ "loss": 0.7831,
346
+ "step": 480
347
+ },
348
+ {
349
+ "epoch": 0.5971968312004875,
350
+ "grad_norm": 0.22364799678325653,
351
+ "learning_rate": 5e-05,
352
+ "loss": 0.7737,
353
+ "step": 490
354
+ },
355
+ {
356
+ "epoch": 0.6093845216331505,
357
+ "grad_norm": 0.23953403532505035,
358
+ "learning_rate": 5e-05,
359
+ "loss": 0.787,
360
+ "step": 500
361
+ },
362
+ {
363
+ "epoch": 0.6215722120658135,
364
+ "grad_norm": 0.22349311411380768,
365
+ "learning_rate": 5e-05,
366
+ "loss": 0.7763,
367
+ "step": 510
368
+ },
369
+ {
370
+ "epoch": 0.6337599024984766,
371
+ "grad_norm": 0.21996377408504486,
372
+ "learning_rate": 5e-05,
373
+ "loss": 0.7596,
374
+ "step": 520
375
+ },
376
+ {
377
+ "epoch": 0.6459475929311396,
378
+ "grad_norm": 0.245538130402565,
379
+ "learning_rate": 5e-05,
380
+ "loss": 0.7785,
381
+ "step": 530
382
+ },
383
+ {
384
+ "epoch": 0.6581352833638026,
385
+ "grad_norm": 0.23442520201206207,
386
+ "learning_rate": 5e-05,
387
+ "loss": 0.7633,
388
+ "step": 540
389
+ },
390
+ {
391
+ "epoch": 0.6703229737964655,
392
+ "grad_norm": 0.22274106740951538,
393
+ "learning_rate": 5e-05,
394
+ "loss": 0.7756,
395
+ "step": 550
396
+ },
397
+ {
398
+ "epoch": 0.6825106642291285,
399
+ "grad_norm": 0.21717530488967896,
400
+ "learning_rate": 5e-05,
401
+ "loss": 0.7726,
402
+ "step": 560
403
+ },
404
+ {
405
+ "epoch": 0.6946983546617916,
406
+ "grad_norm": 0.2465955913066864,
407
+ "learning_rate": 5e-05,
408
+ "loss": 0.7476,
409
+ "step": 570
410
+ },
411
+ {
412
+ "epoch": 0.7068860450944546,
413
+ "grad_norm": 0.2318575084209442,
414
+ "learning_rate": 5e-05,
415
+ "loss": 0.7584,
416
+ "step": 580
417
+ },
418
+ {
419
+ "epoch": 0.7190737355271176,
420
+ "grad_norm": 0.25114378333091736,
421
+ "learning_rate": 5e-05,
422
+ "loss": 0.7625,
423
+ "step": 590
424
+ },
425
+ {
426
+ "epoch": 0.7312614259597806,
427
+ "grad_norm": 0.19145755469799042,
428
+ "learning_rate": 5e-05,
429
+ "loss": 0.7577,
430
+ "step": 600
431
+ },
432
+ {
433
+ "epoch": 0.7434491163924436,
434
+ "grad_norm": 0.19129066169261932,
435
+ "learning_rate": 5e-05,
436
+ "loss": 0.7772,
437
+ "step": 610
438
+ },
439
+ {
440
+ "epoch": 0.7556368068251066,
441
+ "grad_norm": 0.2014475166797638,
442
+ "learning_rate": 5e-05,
443
+ "loss": 0.7601,
444
+ "step": 620
445
+ },
446
+ {
447
+ "epoch": 0.7678244972577697,
448
+ "grad_norm": 0.21480509638786316,
449
+ "learning_rate": 5e-05,
450
+ "loss": 0.7665,
451
+ "step": 630
452
+ },
453
+ {
454
+ "epoch": 0.7800121876904327,
455
+ "grad_norm": 0.2633228302001953,
456
+ "learning_rate": 5e-05,
457
+ "loss": 0.7565,
458
+ "step": 640
459
+ },
460
+ {
461
+ "epoch": 0.7921998781230957,
462
+ "grad_norm": 0.24633203446865082,
463
+ "learning_rate": 5e-05,
464
+ "loss": 0.7602,
465
+ "step": 650
466
+ },
467
+ {
468
+ "epoch": 0.8043875685557587,
469
+ "grad_norm": 0.19770534336566925,
470
+ "learning_rate": 5e-05,
471
+ "loss": 0.7685,
472
+ "step": 660
473
+ },
474
+ {
475
+ "epoch": 0.8165752589884216,
476
+ "grad_norm": 0.2267930805683136,
477
+ "learning_rate": 5e-05,
478
+ "loss": 0.7516,
479
+ "step": 670
480
+ },
481
+ {
482
+ "epoch": 0.8287629494210847,
483
+ "grad_norm": 0.238111674785614,
484
+ "learning_rate": 5e-05,
485
+ "loss": 0.7487,
486
+ "step": 680
487
+ },
488
+ {
489
+ "epoch": 0.8409506398537477,
490
+ "grad_norm": 0.19528904557228088,
491
+ "learning_rate": 5e-05,
492
+ "loss": 0.7467,
493
+ "step": 690
494
+ },
495
+ {
496
+ "epoch": 0.8531383302864107,
497
+ "grad_norm": 0.21305914223194122,
498
+ "learning_rate": 5e-05,
499
+ "loss": 0.7527,
500
+ "step": 700
501
+ },
502
+ {
503
+ "epoch": 0.8653260207190737,
504
+ "grad_norm": 0.2205205112695694,
505
+ "learning_rate": 5e-05,
506
+ "loss": 0.7639,
507
+ "step": 710
508
+ },
509
+ {
510
+ "epoch": 0.8775137111517367,
511
+ "grad_norm": 0.22429127991199493,
512
+ "learning_rate": 5e-05,
513
+ "loss": 0.7744,
514
+ "step": 720
515
+ },
516
+ {
517
+ "epoch": 0.8897014015843998,
518
+ "grad_norm": 0.23755772411823273,
519
+ "learning_rate": 5e-05,
520
+ "loss": 0.7446,
521
+ "step": 730
522
+ },
523
+ {
524
+ "epoch": 0.9018890920170628,
525
+ "grad_norm": 0.17592693865299225,
526
+ "learning_rate": 5e-05,
527
+ "loss": 0.7442,
528
+ "step": 740
529
+ },
530
+ {
531
+ "epoch": 0.9140767824497258,
532
+ "grad_norm": 0.2266898900270462,
533
+ "learning_rate": 5e-05,
534
+ "loss": 0.7326,
535
+ "step": 750
536
+ },
537
+ {
538
+ "epoch": 0.9262644728823888,
539
+ "grad_norm": 0.23541569709777832,
540
+ "learning_rate": 5e-05,
541
+ "loss": 0.7338,
542
+ "step": 760
543
+ },
544
+ {
545
+ "epoch": 0.9384521633150518,
546
+ "grad_norm": 0.22233888506889343,
547
+ "learning_rate": 5e-05,
548
+ "loss": 0.7385,
549
+ "step": 770
550
+ },
551
+ {
552
+ "epoch": 0.9506398537477148,
553
+ "grad_norm": 0.23714767396450043,
554
+ "learning_rate": 5e-05,
555
+ "loss": 0.7306,
556
+ "step": 780
557
+ },
558
+ {
559
+ "epoch": 0.9628275441803779,
560
+ "grad_norm": 0.21191060543060303,
561
+ "learning_rate": 5e-05,
562
+ "loss": 0.736,
563
+ "step": 790
564
+ },
565
+ {
566
+ "epoch": 0.9750152346130408,
567
+ "grad_norm": 0.22759106755256653,
568
+ "learning_rate": 5e-05,
569
+ "loss": 0.7505,
570
+ "step": 800
571
+ },
572
+ {
573
+ "epoch": 0.9872029250457038,
574
+ "grad_norm": 0.19895903766155243,
575
+ "learning_rate": 5e-05,
576
+ "loss": 0.7536,
577
+ "step": 810
578
+ },
579
+ {
580
+ "epoch": 0.9993906154783668,
581
+ "grad_norm": 0.3061685860157013,
582
+ "learning_rate": 5e-05,
583
+ "loss": 0.7493,
584
+ "step": 820
585
+ },
586
+ {
587
+ "epoch": 1.0109689213893966,
588
+ "grad_norm": 0.18868015706539154,
589
+ "learning_rate": 5e-05,
590
+ "loss": 0.7342,
591
+ "step": 830
592
+ },
593
+ {
594
+ "epoch": 1.0231566118220596,
595
+ "grad_norm": 0.1797628402709961,
596
+ "learning_rate": 5e-05,
597
+ "loss": 0.7192,
598
+ "step": 840
599
+ },
600
+ {
601
+ "epoch": 1.0353443022547226,
602
+ "grad_norm": 0.19511596858501434,
603
+ "learning_rate": 5e-05,
604
+ "loss": 0.7002,
605
+ "step": 850
606
+ },
607
+ {
608
+ "epoch": 1.0475319926873858,
609
+ "grad_norm": 0.2123185396194458,
610
+ "learning_rate": 5e-05,
611
+ "loss": 0.7236,
612
+ "step": 860
613
+ },
614
+ {
615
+ "epoch": 1.0597196831200488,
616
+ "grad_norm": 0.2047470211982727,
617
+ "learning_rate": 5e-05,
618
+ "loss": 0.711,
619
+ "step": 870
620
+ },
621
+ {
622
+ "epoch": 1.0719073735527118,
623
+ "grad_norm": 0.18809868395328522,
624
+ "learning_rate": 5e-05,
625
+ "loss": 0.7253,
626
+ "step": 880
627
+ },
628
+ {
629
+ "epoch": 1.0840950639853748,
630
+ "grad_norm": 0.2207012176513672,
631
+ "learning_rate": 5e-05,
632
+ "loss": 0.6971,
633
+ "step": 890
634
+ },
635
+ {
636
+ "epoch": 1.0962827544180378,
637
+ "grad_norm": 0.1912035495042801,
638
+ "learning_rate": 5e-05,
639
+ "loss": 0.7239,
640
+ "step": 900
641
+ },
642
+ {
643
+ "epoch": 1.1084704448507008,
644
+ "grad_norm": 0.23486360907554626,
645
+ "learning_rate": 5e-05,
646
+ "loss": 0.7335,
647
+ "step": 910
648
+ },
649
+ {
650
+ "epoch": 1.1206581352833638,
651
+ "grad_norm": 0.22189787030220032,
652
+ "learning_rate": 5e-05,
653
+ "loss": 0.7149,
654
+ "step": 920
655
+ },
656
+ {
657
+ "epoch": 1.1328458257160268,
658
+ "grad_norm": 0.25818824768066406,
659
+ "learning_rate": 5e-05,
660
+ "loss": 0.7322,
661
+ "step": 930
662
+ },
663
+ {
664
+ "epoch": 1.1450335161486898,
665
+ "grad_norm": 0.2088775783777237,
666
+ "learning_rate": 5e-05,
667
+ "loss": 0.7382,
668
+ "step": 940
669
+ },
670
+ {
671
+ "epoch": 1.1572212065813527,
672
+ "grad_norm": 0.19445833563804626,
673
+ "learning_rate": 5e-05,
674
+ "loss": 0.7223,
675
+ "step": 950
676
+ },
677
+ {
678
+ "epoch": 1.1694088970140157,
679
+ "grad_norm": 0.21489456295967102,
680
+ "learning_rate": 5e-05,
681
+ "loss": 0.7317,
682
+ "step": 960
683
+ },
684
+ {
685
+ "epoch": 1.181596587446679,
686
+ "grad_norm": 0.24050794541835785,
687
+ "learning_rate": 5e-05,
688
+ "loss": 0.7275,
689
+ "step": 970
690
+ },
691
+ {
692
+ "epoch": 1.193784277879342,
693
+ "grad_norm": 0.18142102658748627,
694
+ "learning_rate": 5e-05,
695
+ "loss": 0.7046,
696
+ "step": 980
697
+ },
698
+ {
699
+ "epoch": 1.205971968312005,
700
+ "grad_norm": 0.4781077802181244,
701
+ "learning_rate": 5e-05,
702
+ "loss": 0.7433,
703
+ "step": 990
704
+ },
705
+ {
706
+ "epoch": 1.218159658744668,
707
+ "grad_norm": 0.20628023147583008,
708
+ "learning_rate": 5e-05,
709
+ "loss": 0.7403,
710
+ "step": 1000
711
+ },
712
+ {
713
+ "epoch": 1.230347349177331,
714
+ "grad_norm": 0.19117140769958496,
715
+ "learning_rate": 5e-05,
716
+ "loss": 0.7308,
717
+ "step": 1010
718
+ },
719
+ {
720
+ "epoch": 1.242535039609994,
721
+ "grad_norm": 0.2107745110988617,
722
+ "learning_rate": 5e-05,
723
+ "loss": 0.7201,
724
+ "step": 1020
725
+ },
726
+ {
727
+ "epoch": 1.2547227300426569,
728
+ "grad_norm": 0.20551154017448425,
729
+ "learning_rate": 5e-05,
730
+ "loss": 0.7038,
731
+ "step": 1030
732
+ },
733
+ {
734
+ "epoch": 1.2669104204753199,
735
+ "grad_norm": 0.26577022671699524,
736
+ "learning_rate": 5e-05,
737
+ "loss": 0.7096,
738
+ "step": 1040
739
+ },
740
+ {
741
+ "epoch": 1.2790981109079829,
742
+ "grad_norm": 0.22094209492206573,
743
+ "learning_rate": 5e-05,
744
+ "loss": 0.718,
745
+ "step": 1050
746
+ },
747
+ {
748
+ "epoch": 1.291285801340646,
749
+ "grad_norm": 0.18224076926708221,
750
+ "learning_rate": 5e-05,
751
+ "loss": 0.7261,
752
+ "step": 1060
753
+ },
754
+ {
755
+ "epoch": 1.303473491773309,
756
+ "grad_norm": 0.226251482963562,
757
+ "learning_rate": 5e-05,
758
+ "loss": 0.7372,
759
+ "step": 1070
760
+ },
761
+ {
762
+ "epoch": 1.315661182205972,
763
+ "grad_norm": 0.20186461508274078,
764
+ "learning_rate": 5e-05,
765
+ "loss": 0.7074,
766
+ "step": 1080
767
+ },
768
+ {
769
+ "epoch": 1.327848872638635,
770
+ "grad_norm": 0.22229762375354767,
771
+ "learning_rate": 5e-05,
772
+ "loss": 0.7069,
773
+ "step": 1090
774
+ },
775
+ {
776
+ "epoch": 1.340036563071298,
777
+ "grad_norm": 0.1776452213525772,
778
+ "learning_rate": 5e-05,
779
+ "loss": 0.714,
780
+ "step": 1100
781
+ },
782
+ {
783
+ "epoch": 1.352224253503961,
784
+ "grad_norm": 0.20680229365825653,
785
+ "learning_rate": 5e-05,
786
+ "loss": 0.7066,
787
+ "step": 1110
788
+ },
789
+ {
790
+ "epoch": 1.364411943936624,
791
+ "grad_norm": 0.21842031180858612,
792
+ "learning_rate": 5e-05,
793
+ "loss": 0.7281,
794
+ "step": 1120
795
+ },
796
+ {
797
+ "epoch": 1.376599634369287,
798
+ "grad_norm": 0.28037774562835693,
799
+ "learning_rate": 5e-05,
800
+ "loss": 0.7149,
801
+ "step": 1130
802
+ },
803
+ {
804
+ "epoch": 1.38878732480195,
805
+ "grad_norm": 0.20962920784950256,
806
+ "learning_rate": 5e-05,
807
+ "loss": 0.6951,
808
+ "step": 1140
809
+ },
810
+ {
811
+ "epoch": 1.400975015234613,
812
+ "grad_norm": 0.18751871585845947,
813
+ "learning_rate": 5e-05,
814
+ "loss": 0.7101,
815
+ "step": 1150
816
+ },
817
+ {
818
+ "epoch": 1.413162705667276,
819
+ "grad_norm": 0.2722964584827423,
820
+ "learning_rate": 5e-05,
821
+ "loss": 0.7205,
822
+ "step": 1160
823
+ },
824
+ {
825
+ "epoch": 1.425350396099939,
826
+ "grad_norm": 0.20932041108608246,
827
+ "learning_rate": 5e-05,
828
+ "loss": 0.7082,
829
+ "step": 1170
830
+ },
831
+ {
832
+ "epoch": 1.437538086532602,
833
+ "grad_norm": 0.17836257815361023,
834
+ "learning_rate": 5e-05,
835
+ "loss": 0.7232,
836
+ "step": 1180
837
+ },
838
+ {
839
+ "epoch": 1.449725776965265,
840
+ "grad_norm": 0.3180703818798065,
841
+ "learning_rate": 5e-05,
842
+ "loss": 0.7245,
843
+ "step": 1190
844
+ },
845
+ {
846
+ "epoch": 1.4619134673979282,
847
+ "grad_norm": 0.201424703001976,
848
+ "learning_rate": 5e-05,
849
+ "loss": 0.7232,
850
+ "step": 1200
851
+ },
852
+ {
853
+ "epoch": 1.4741011578305911,
854
+ "grad_norm": 0.20228999853134155,
855
+ "learning_rate": 5e-05,
856
+ "loss": 0.7169,
857
+ "step": 1210
858
+ },
859
+ {
860
+ "epoch": 1.4862888482632541,
861
+ "grad_norm": 0.17478394508361816,
862
+ "learning_rate": 5e-05,
863
+ "loss": 0.704,
864
+ "step": 1220
865
+ },
866
+ {
867
+ "epoch": 1.4984765386959171,
868
+ "grad_norm": 0.20061351358890533,
869
+ "learning_rate": 5e-05,
870
+ "loss": 0.6792,
871
+ "step": 1230
872
+ },
873
+ {
874
+ "epoch": 1.5106642291285801,
875
+ "grad_norm": 0.1857377588748932,
876
+ "learning_rate": 5e-05,
877
+ "loss": 0.709,
878
+ "step": 1240
879
+ },
880
+ {
881
+ "epoch": 1.522851919561243,
882
+ "grad_norm": 0.19873379170894623,
883
+ "learning_rate": 5e-05,
884
+ "loss": 0.7007,
885
+ "step": 1250
886
+ },
887
+ {
888
+ "epoch": 1.5350396099939063,
889
+ "grad_norm": 0.19807074964046478,
890
+ "learning_rate": 5e-05,
891
+ "loss": 0.7122,
892
+ "step": 1260
893
+ },
894
+ {
895
+ "epoch": 1.5472273004265693,
896
+ "grad_norm": 0.16746796667575836,
897
+ "learning_rate": 5e-05,
898
+ "loss": 0.7005,
899
+ "step": 1270
900
+ },
901
+ {
902
+ "epoch": 1.5594149908592323,
903
+ "grad_norm": 0.2750326991081238,
904
+ "learning_rate": 5e-05,
905
+ "loss": 0.7166,
906
+ "step": 1280
907
+ },
908
+ {
909
+ "epoch": 1.5716026812918953,
910
+ "grad_norm": 0.2776738107204437,
911
+ "learning_rate": 5e-05,
912
+ "loss": 0.7405,
913
+ "step": 1290
914
+ },
915
+ {
916
+ "epoch": 1.5837903717245583,
917
+ "grad_norm": 0.3494517207145691,
918
+ "learning_rate": 5e-05,
919
+ "loss": 0.7192,
920
+ "step": 1300
921
+ },
922
+ {
923
+ "epoch": 1.5959780621572213,
924
+ "grad_norm": 0.25028494000434875,
925
+ "learning_rate": 5e-05,
926
+ "loss": 0.741,
927
+ "step": 1310
928
+ },
929
+ {
930
+ "epoch": 1.6081657525898843,
931
+ "grad_norm": 0.23914143443107605,
932
+ "learning_rate": 5e-05,
933
+ "loss": 0.72,
934
+ "step": 1320
935
+ },
936
+ {
937
+ "epoch": 1.6203534430225472,
938
+ "grad_norm": 0.2052862048149109,
939
+ "learning_rate": 5e-05,
940
+ "loss": 0.7112,
941
+ "step": 1330
942
+ },
943
+ {
944
+ "epoch": 1.6325411334552102,
945
+ "grad_norm": 0.22285936772823334,
946
+ "learning_rate": 5e-05,
947
+ "loss": 0.7165,
948
+ "step": 1340
949
+ },
950
+ {
951
+ "epoch": 1.6447288238878732,
952
+ "grad_norm": 0.20810848474502563,
953
+ "learning_rate": 5e-05,
954
+ "loss": 0.7122,
955
+ "step": 1350
956
+ },
957
+ {
958
+ "epoch": 1.6569165143205362,
959
+ "grad_norm": 0.20422884821891785,
960
+ "learning_rate": 5e-05,
961
+ "loss": 0.7082,
962
+ "step": 1360
963
+ },
964
+ {
965
+ "epoch": 1.6691042047531992,
966
+ "grad_norm": 0.18164314329624176,
967
+ "learning_rate": 5e-05,
968
+ "loss": 0.7131,
969
+ "step": 1370
970
+ },
971
+ {
972
+ "epoch": 1.6812918951858622,
973
+ "grad_norm": 0.19893014430999756,
974
+ "learning_rate": 5e-05,
975
+ "loss": 0.7055,
976
+ "step": 1380
977
+ },
978
+ {
979
+ "epoch": 1.6934795856185252,
980
+ "grad_norm": 0.24865971505641937,
981
+ "learning_rate": 5e-05,
982
+ "loss": 0.6994,
983
+ "step": 1390
984
+ },
985
+ {
986
+ "epoch": 1.7056672760511882,
987
+ "grad_norm": 0.18078123033046722,
988
+ "learning_rate": 5e-05,
989
+ "loss": 0.7162,
990
+ "step": 1400
991
+ },
992
+ {
993
+ "epoch": 1.7178549664838512,
994
+ "grad_norm": 0.1973421722650528,
995
+ "learning_rate": 5e-05,
996
+ "loss": 0.7136,
997
+ "step": 1410
998
+ },
999
+ {
1000
+ "epoch": 1.7300426569165142,
1001
+ "grad_norm": 0.2057911604642868,
1002
+ "learning_rate": 5e-05,
1003
+ "loss": 0.7004,
1004
+ "step": 1420
1005
+ },
1006
+ {
1007
+ "epoch": 1.7422303473491774,
1008
+ "grad_norm": 0.18640495836734772,
1009
+ "learning_rate": 5e-05,
1010
+ "loss": 0.7102,
1011
+ "step": 1430
1012
+ },
1013
+ {
1014
+ "epoch": 1.7544180377818404,
1015
+ "grad_norm": 0.21038417518138885,
1016
+ "learning_rate": 5e-05,
1017
+ "loss": 0.6908,
1018
+ "step": 1440
1019
+ },
1020
+ {
1021
+ "epoch": 1.7666057282145033,
1022
+ "grad_norm": 0.2123890221118927,
1023
+ "learning_rate": 5e-05,
1024
+ "loss": 0.6913,
1025
+ "step": 1450
1026
+ },
1027
+ {
1028
+ "epoch": 1.7787934186471663,
1029
+ "grad_norm": 0.21619991958141327,
1030
+ "learning_rate": 5e-05,
1031
+ "loss": 0.7038,
1032
+ "step": 1460
1033
+ },
1034
+ {
1035
+ "epoch": 1.7909811090798293,
1036
+ "grad_norm": 0.1975220888853073,
1037
+ "learning_rate": 5e-05,
1038
+ "loss": 0.7041,
1039
+ "step": 1470
1040
+ },
1041
+ {
1042
+ "epoch": 1.8031687995124925,
1043
+ "grad_norm": 0.20595207810401917,
1044
+ "learning_rate": 5e-05,
1045
+ "loss": 0.7145,
1046
+ "step": 1480
1047
+ },
1048
+ {
1049
+ "epoch": 1.8153564899451555,
1050
+ "grad_norm": 0.16633334755897522,
1051
+ "learning_rate": 5e-05,
1052
+ "loss": 0.6963,
1053
+ "step": 1490
1054
+ },
1055
+ {
1056
+ "epoch": 1.8275441803778185,
1057
+ "grad_norm": 0.19140620529651642,
1058
+ "learning_rate": 5e-05,
1059
+ "loss": 0.6814,
1060
+ "step": 1500
1061
+ },
1062
+ {
1063
+ "epoch": 1.8397318708104815,
1064
+ "grad_norm": 0.1833905130624771,
1065
+ "learning_rate": 5e-05,
1066
+ "loss": 0.7174,
1067
+ "step": 1510
1068
+ },
1069
+ {
1070
+ "epoch": 1.8519195612431445,
1071
+ "grad_norm": 0.20587337017059326,
1072
+ "learning_rate": 5e-05,
1073
+ "loss": 0.7098,
1074
+ "step": 1520
1075
+ },
1076
+ {
1077
+ "epoch": 1.8641072516758075,
1078
+ "grad_norm": 0.21952909231185913,
1079
+ "learning_rate": 5e-05,
1080
+ "loss": 0.6912,
1081
+ "step": 1530
1082
+ },
1083
+ {
1084
+ "epoch": 1.8762949421084705,
1085
+ "grad_norm": 0.23233620822429657,
1086
+ "learning_rate": 5e-05,
1087
+ "loss": 0.6992,
1088
+ "step": 1540
1089
+ },
1090
+ {
1091
+ "epoch": 1.8884826325411335,
1092
+ "grad_norm": 0.17159809172153473,
1093
+ "learning_rate": 5e-05,
1094
+ "loss": 0.7052,
1095
+ "step": 1550
1096
+ },
1097
+ {
1098
+ "epoch": 1.9006703229737965,
1099
+ "grad_norm": 0.23522134125232697,
1100
+ "learning_rate": 5e-05,
1101
+ "loss": 0.6873,
1102
+ "step": 1560
1103
+ },
1104
+ {
1105
+ "epoch": 1.9128580134064594,
1106
+ "grad_norm": 0.21959669888019562,
1107
+ "learning_rate": 5e-05,
1108
+ "loss": 0.6834,
1109
+ "step": 1570
1110
+ },
1111
+ {
1112
+ "epoch": 1.9250457038391224,
1113
+ "grad_norm": 0.2110578864812851,
1114
+ "learning_rate": 5e-05,
1115
+ "loss": 0.7038,
1116
+ "step": 1580
1117
+ },
1118
+ {
1119
+ "epoch": 1.9372333942717854,
1120
+ "grad_norm": 0.18674372136592865,
1121
+ "learning_rate": 5e-05,
1122
+ "loss": 0.6743,
1123
+ "step": 1590
1124
+ },
1125
+ {
1126
+ "epoch": 1.9494210847044484,
1127
+ "grad_norm": 0.18220046162605286,
1128
+ "learning_rate": 5e-05,
1129
+ "loss": 0.6943,
1130
+ "step": 1600
1131
+ },
1132
+ {
1133
+ "epoch": 1.9616087751371114,
1134
+ "grad_norm": 0.20756569504737854,
1135
+ "learning_rate": 5e-05,
1136
+ "loss": 0.7049,
1137
+ "step": 1610
1138
+ },
1139
+ {
1140
+ "epoch": 1.9737964655697744,
1141
+ "grad_norm": 0.1925211250782013,
1142
+ "learning_rate": 5e-05,
1143
+ "loss": 0.7033,
1144
+ "step": 1620
1145
+ },
1146
+ {
1147
+ "epoch": 1.9859841560024374,
1148
+ "grad_norm": 0.21374447643756866,
1149
+ "learning_rate": 5e-05,
1150
+ "loss": 0.7025,
1151
+ "step": 1630
1152
+ },
1153
+ {
1154
+ "epoch": 1.9981718464351006,
1155
+ "grad_norm": 0.1952880322933197,
1156
+ "learning_rate": 5e-05,
1157
+ "loss": 0.6877,
1158
+ "step": 1640
1159
+ },
1160
+ {
1161
+ "epoch": 2.0097501523461303,
1162
+ "grad_norm": 0.20457723736763,
1163
+ "learning_rate": 5e-05,
1164
+ "loss": 0.6734,
1165
+ "step": 1650
1166
+ },
1167
+ {
1168
+ "epoch": 2.0219378427787933,
1169
+ "grad_norm": 0.20965932309627533,
1170
+ "learning_rate": 5e-05,
1171
+ "loss": 0.6472,
1172
+ "step": 1660
1173
+ },
1174
+ {
1175
+ "epoch": 2.0341255332114563,
1176
+ "grad_norm": 0.19478890299797058,
1177
+ "learning_rate": 5e-05,
1178
+ "loss": 0.6781,
1179
+ "step": 1670
1180
+ },
1181
+ {
1182
+ "epoch": 2.0463132236441193,
1183
+ "grad_norm": 0.21520769596099854,
1184
+ "learning_rate": 5e-05,
1185
+ "loss": 0.6592,
1186
+ "step": 1680
1187
+ },
1188
+ {
1189
+ "epoch": 2.0585009140767823,
1190
+ "grad_norm": 0.20154155790805817,
1191
+ "learning_rate": 5e-05,
1192
+ "loss": 0.6744,
1193
+ "step": 1690
1194
+ },
1195
+ {
1196
+ "epoch": 2.0706886045094453,
1197
+ "grad_norm": 0.1899706870317459,
1198
+ "learning_rate": 5e-05,
1199
+ "loss": 0.6592,
1200
+ "step": 1700
1201
+ },
1202
+ {
1203
+ "epoch": 2.0828762949421087,
1204
+ "grad_norm": 0.20379123091697693,
1205
+ "learning_rate": 5e-05,
1206
+ "loss": 0.6688,
1207
+ "step": 1710
1208
+ },
1209
+ {
1210
+ "epoch": 2.0950639853747717,
1211
+ "grad_norm": 0.19799865782260895,
1212
+ "learning_rate": 5e-05,
1213
+ "loss": 0.6678,
1214
+ "step": 1720
1215
+ },
1216
+ {
1217
+ "epoch": 2.1072516758074347,
1218
+ "grad_norm": 0.19186025857925415,
1219
+ "learning_rate": 5e-05,
1220
+ "loss": 0.6587,
1221
+ "step": 1730
1222
+ },
1223
+ {
1224
+ "epoch": 2.1194393662400977,
1225
+ "grad_norm": 0.18936173617839813,
1226
+ "learning_rate": 5e-05,
1227
+ "loss": 0.668,
1228
+ "step": 1740
1229
+ },
1230
+ {
1231
+ "epoch": 2.1316270566727606,
1232
+ "grad_norm": 0.19159527122974396,
1233
+ "learning_rate": 5e-05,
1234
+ "loss": 0.6643,
1235
+ "step": 1750
1236
+ },
1237
+ {
1238
+ "epoch": 2.1438147471054236,
1239
+ "grad_norm": 0.20596368610858917,
1240
+ "learning_rate": 5e-05,
1241
+ "loss": 0.6667,
1242
+ "step": 1760
1243
+ },
1244
+ {
1245
+ "epoch": 2.1560024375380866,
1246
+ "grad_norm": 0.18601712584495544,
1247
+ "learning_rate": 5e-05,
1248
+ "loss": 0.6598,
1249
+ "step": 1770
1250
+ },
1251
+ {
1252
+ "epoch": 2.1681901279707496,
1253
+ "grad_norm": 0.2130655199289322,
1254
+ "learning_rate": 5e-05,
1255
+ "loss": 0.6844,
1256
+ "step": 1780
1257
+ },
1258
+ {
1259
+ "epoch": 2.1803778184034126,
1260
+ "grad_norm": 0.20477338135242462,
1261
+ "learning_rate": 5e-05,
1262
+ "loss": 0.666,
1263
+ "step": 1790
1264
+ },
1265
+ {
1266
+ "epoch": 2.1925655088360756,
1267
+ "grad_norm": 0.18484896421432495,
1268
+ "learning_rate": 5e-05,
1269
+ "loss": 0.6659,
1270
+ "step": 1800
1271
+ },
1272
+ {
1273
+ "epoch": 2.2047531992687386,
1274
+ "grad_norm": 0.17023979127407074,
1275
+ "learning_rate": 5e-05,
1276
+ "loss": 0.6721,
1277
+ "step": 1810
1278
+ },
1279
+ {
1280
+ "epoch": 2.2169408897014016,
1281
+ "grad_norm": 0.23462794721126556,
1282
+ "learning_rate": 5e-05,
1283
+ "loss": 0.6699,
1284
+ "step": 1820
1285
+ },
1286
+ {
1287
+ "epoch": 2.2291285801340646,
1288
+ "grad_norm": 0.16521663963794708,
1289
+ "learning_rate": 5e-05,
1290
+ "loss": 0.6584,
1291
+ "step": 1830
1292
+ },
1293
+ {
1294
+ "epoch": 2.2413162705667276,
1295
+ "grad_norm": 0.19637759029865265,
1296
+ "learning_rate": 5e-05,
1297
+ "loss": 0.6621,
1298
+ "step": 1840
1299
+ }
1300
+ ],
1301
+ "logging_steps": 10,
1302
+ "max_steps": 2463,
1303
+ "num_input_tokens_seen": 0,
1304
+ "num_train_epochs": 3,
1305
+ "save_steps": 205,
1306
+ "stateful_callbacks": {
1307
+ "TrainerControl": {
1308
+ "args": {
1309
+ "should_epoch_stop": false,
1310
+ "should_evaluate": false,
1311
+ "should_log": false,
1312
+ "should_save": true,
1313
+ "should_training_stop": false
1314
+ },
1315
+ "attributes": {}
1316
+ }
1317
+ },
1318
+ "total_flos": 7.565604904304116e+18,
1319
+ "train_batch_size": 4,
1320
+ "trial_name": null,
1321
+ "trial_params": null
1322
+ }
openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1845/training_args.bin ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:56896ea21940e33d7e2934267dff2a5b4a4a80195a0f165245f6617a0f88b90b
3
+ size 7377
openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-1845/zero_to_fp32.py ADDED
@@ -0,0 +1,760 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python
2
+
3
+ # Copyright (c) Microsoft Corporation.
4
+ # SPDX-License-Identifier: Apache-2.0
5
+
6
+ # DeepSpeed Team
7
+
8
+ # This script extracts fp32 consolidated weights from a zero 1, 2 and 3 DeepSpeed checkpoints. It gets
9
+ # copied into the top level checkpoint dir, so the user can easily do the conversion at any point in
10
+ # the future. Once extracted, the weights don't require DeepSpeed and can be used in any
11
+ # application.
12
+ #
13
+ # example:
14
+ # python zero_to_fp32.py . output_dir/
15
+ # or
16
+ # python zero_to_fp32.py . output_dir/ --safe_serialization
17
+
18
+ import argparse
19
+ import torch
20
+ import glob
21
+ import math
22
+ import os
23
+ import re
24
+ import gc
25
+ import json
26
+ import numpy as np
27
+ from tqdm import tqdm
28
+ from collections import OrderedDict
29
+ from dataclasses import dataclass
30
+
31
+ # while this script doesn't use deepspeed to recover data, since the checkpoints are pickled with
32
+ # DeepSpeed data structures it has to be available in the current python environment.
33
+ from deepspeed.utils import logger
34
+ from deepspeed.checkpoint.constants import (DS_VERSION, OPTIMIZER_STATE_DICT, SINGLE_PARTITION_OF_FP32_GROUPS,
35
+ FP32_FLAT_GROUPS, ZERO_STAGE, PARTITION_COUNT, PARAM_SHAPES, BUFFER_NAMES,
36
+ FROZEN_PARAM_SHAPES, FROZEN_PARAM_FRAGMENTS)
37
+
38
+
39
+ @dataclass
40
+ class zero_model_state:
41
+ buffers: dict()
42
+ param_shapes: dict()
43
+ shared_params: list
44
+ ds_version: int
45
+ frozen_param_shapes: dict()
46
+ frozen_param_fragments: dict()
47
+
48
+
49
+ debug = 0
50
+
51
+ # load to cpu
52
+ device = torch.device('cpu')
53
+
54
+
55
+ def atoi(text):
56
+ return int(text) if text.isdigit() else text
57
+
58
+
59
+ def natural_keys(text):
60
+ '''
61
+ alist.sort(key=natural_keys) sorts in human order
62
+ http://nedbatchelder.com/blog/200712/human_sorting.html
63
+ (See Toothy's implementation in the comments)
64
+ '''
65
+ return [atoi(c) for c in re.split(r'(\d+)', text)]
66
+
67
+
68
+ def get_model_state_file(checkpoint_dir, zero_stage):
69
+ if not os.path.isdir(checkpoint_dir):
70
+ raise FileNotFoundError(f"Directory '{checkpoint_dir}' doesn't exist")
71
+
72
+ # there should be only one file
73
+ if zero_stage <= 2:
74
+ file = os.path.join(checkpoint_dir, "mp_rank_00_model_states.pt")
75
+ elif zero_stage == 3:
76
+ file = os.path.join(checkpoint_dir, "zero_pp_rank_0_mp_rank_00_model_states.pt")
77
+
78
+ if not os.path.exists(file):
79
+ raise FileNotFoundError(f"can't find model states file at '{file}'")
80
+
81
+ return file
82
+
83
+
84
+ def get_checkpoint_files(checkpoint_dir, glob_pattern):
85
+ # XXX: need to test that this simple glob rule works for multi-node setup too
86
+ ckpt_files = sorted(glob.glob(os.path.join(checkpoint_dir, glob_pattern)), key=natural_keys)
87
+
88
+ if len(ckpt_files) == 0:
89
+ raise FileNotFoundError(f"can't find {glob_pattern} files in directory '{checkpoint_dir}'")
90
+
91
+ return ckpt_files
92
+
93
+
94
+ def get_optim_files(checkpoint_dir):
95
+ return get_checkpoint_files(checkpoint_dir, "*_optim_states.pt")
96
+
97
+
98
+ def get_model_state_files(checkpoint_dir):
99
+ return get_checkpoint_files(checkpoint_dir, "*_model_states.pt")
100
+
101
+
102
+ def parse_model_states(files):
103
+ zero_model_states = []
104
+ for file in files:
105
+ state_dict = torch.load(file, map_location=device, weights_only=False)
106
+
107
+ if BUFFER_NAMES not in state_dict:
108
+ raise ValueError(f"{file} is not a model state checkpoint")
109
+ buffer_names = state_dict[BUFFER_NAMES]
110
+ if debug:
111
+ print("Found buffers:", buffer_names)
112
+
113
+ # recover just the buffers while restoring them to fp32 if they were saved in fp16
114
+ buffers = {k: v.float() for k, v in state_dict["module"].items() if k in buffer_names}
115
+ param_shapes = state_dict[PARAM_SHAPES]
116
+
117
+ # collect parameters that are included in param_shapes
118
+ param_names = []
119
+ for s in param_shapes:
120
+ for name in s.keys():
121
+ param_names.append(name)
122
+
123
+ # update with frozen parameters
124
+ frozen_param_shapes = state_dict.get(FROZEN_PARAM_SHAPES, None)
125
+ if frozen_param_shapes is not None:
126
+ if debug:
127
+ print(f"Found frozen_param_shapes: {frozen_param_shapes}")
128
+ param_names += list(frozen_param_shapes.keys())
129
+
130
+ # handle shared params
131
+ shared_params = [[k, v] for k, v in state_dict["shared_params"].items()]
132
+
133
+ ds_version = state_dict.get(DS_VERSION, None)
134
+
135
+ frozen_param_fragments = state_dict.get(FROZEN_PARAM_FRAGMENTS, None)
136
+
137
+ z_model_state = zero_model_state(buffers=buffers,
138
+ param_shapes=param_shapes,
139
+ shared_params=shared_params,
140
+ ds_version=ds_version,
141
+ frozen_param_shapes=frozen_param_shapes,
142
+ frozen_param_fragments=frozen_param_fragments)
143
+ zero_model_states.append(z_model_state)
144
+
145
+ return zero_model_states
146
+
147
+
148
+ def parse_optim_states(files, ds_checkpoint_dir):
149
+ total_files = len(files)
150
+ state_dicts = []
151
+ for f in tqdm(files, desc='Loading checkpoint shards'):
152
+ state_dict = torch.load(f, map_location=device, mmap=True, weights_only=False)
153
+ # immediately discard the potentially huge 2 optimizer states as we only care for fp32 master weights
154
+ # and also handle the case where it was already removed by another helper script
155
+ state_dict["optimizer_state_dict"].pop("optimizer_state_dict", None)
156
+ state_dicts.append(state_dict)
157
+
158
+ if ZERO_STAGE not in state_dicts[0][OPTIMIZER_STATE_DICT]:
159
+ raise ValueError(f"{files[0]} is not a zero checkpoint")
160
+ zero_stage = state_dicts[0][OPTIMIZER_STATE_DICT][ZERO_STAGE]
161
+ world_size = state_dicts[0][OPTIMIZER_STATE_DICT][PARTITION_COUNT]
162
+
163
+ # For ZeRO-2 each param group can have different partition_count as data parallelism for expert
164
+ # parameters can be different from data parallelism for non-expert parameters. So we can just
165
+ # use the max of the partition_count to get the dp world_size.
166
+
167
+ if type(world_size) is list:
168
+ world_size = max(world_size)
169
+
170
+ if world_size != total_files:
171
+ raise ValueError(
172
+ f"Expected {world_size} of '*_optim_states.pt' under '{ds_checkpoint_dir}' but found {total_files} files. "
173
+ "Possibly due to an overwrite of an old checkpoint, or a checkpoint didn't get saved by one or more processes."
174
+ )
175
+
176
+ # the groups are named differently in each stage
177
+ if zero_stage <= 2:
178
+ fp32_groups_key = SINGLE_PARTITION_OF_FP32_GROUPS
179
+ elif zero_stage == 3:
180
+ fp32_groups_key = FP32_FLAT_GROUPS
181
+ else:
182
+ raise ValueError(f"unknown zero stage {zero_stage}")
183
+
184
+ fp32_flat_groups = [state_dicts[i][OPTIMIZER_STATE_DICT][fp32_groups_key] for i in range(len(state_dicts))]
185
+ return zero_stage, world_size, fp32_flat_groups
186
+
187
+
188
+ def _get_fp32_state_dict_from_zero_checkpoint(ds_checkpoint_dir, exclude_frozen_parameters):
189
+ """
190
+ Returns fp32 state_dict reconstructed from ds checkpoint
191
+
192
+ Args:
193
+ - ``ds_checkpoint_dir``: path to the deepspeed checkpoint folder (where the optimizer files are)
194
+
195
+ """
196
+ print(f"Processing zero checkpoint '{ds_checkpoint_dir}'")
197
+
198
+ optim_files = get_optim_files(ds_checkpoint_dir)
199
+ zero_stage, world_size, fp32_flat_groups = parse_optim_states(optim_files, ds_checkpoint_dir)
200
+ print(f"Detected checkpoint of type zero stage {zero_stage}, world_size: {world_size}")
201
+
202
+ model_files = get_model_state_files(ds_checkpoint_dir)
203
+
204
+ zero_model_states = parse_model_states(model_files)
205
+ print(f'Parsing checkpoint created by deepspeed=={zero_model_states[0].ds_version}')
206
+
207
+ if zero_stage <= 2:
208
+ return _get_fp32_state_dict_from_zero2_checkpoint(world_size, fp32_flat_groups, zero_model_states,
209
+ exclude_frozen_parameters)
210
+ elif zero_stage == 3:
211
+ return _get_fp32_state_dict_from_zero3_checkpoint(world_size, fp32_flat_groups, zero_model_states,
212
+ exclude_frozen_parameters)
213
+
214
+
215
+ def _zero2_merge_frozen_params(state_dict, zero_model_states):
216
+ if zero_model_states[0].frozen_param_shapes is None or len(zero_model_states[0].frozen_param_shapes) == 0:
217
+ return
218
+
219
+ frozen_param_shapes = zero_model_states[0].frozen_param_shapes
220
+ frozen_param_fragments = zero_model_states[0].frozen_param_fragments
221
+
222
+ if debug:
223
+ num_elem = sum(s.numel() for s in frozen_param_shapes.values())
224
+ print(f'rank 0: {FROZEN_PARAM_SHAPES}.numel = {num_elem}')
225
+
226
+ wanted_params = len(frozen_param_shapes)
227
+ wanted_numel = sum(s.numel() for s in frozen_param_shapes.values())
228
+ avail_numel = sum([p.numel() for p in frozen_param_fragments.values()])
229
+ print(f'Frozen params: Have {avail_numel} numels to process.')
230
+ print(f'Frozen params: Need {wanted_numel} numels in {wanted_params} params')
231
+
232
+ total_params = 0
233
+ total_numel = 0
234
+ for name, shape in frozen_param_shapes.items():
235
+ total_params += 1
236
+ unpartitioned_numel = shape.numel()
237
+ total_numel += unpartitioned_numel
238
+
239
+ state_dict[name] = frozen_param_fragments[name]
240
+
241
+ if debug:
242
+ print(f"{name} full shape: {shape} unpartitioned numel {unpartitioned_numel} ")
243
+
244
+ print(f"Reconstructed Frozen fp32 state dict with {total_params} params {total_numel} elements")
245
+
246
+
247
+ def _has_callable(obj, fn):
248
+ attr = getattr(obj, fn, None)
249
+ return callable(attr)
250
+
251
+
252
+ def _zero2_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states):
253
+ param_shapes = zero_model_states[0].param_shapes
254
+
255
+ # Reconstruction protocol:
256
+ #
257
+ # XXX: document this
258
+
259
+ if debug:
260
+ for i in range(world_size):
261
+ for j in range(len(fp32_flat_groups[0])):
262
+ print(f"{FP32_FLAT_GROUPS}[{i}][{j}].shape={fp32_flat_groups[i][j].shape}")
263
+
264
+ # XXX: memory usage doubles here (zero2)
265
+ num_param_groups = len(fp32_flat_groups[0])
266
+ merged_single_partition_of_fp32_groups = []
267
+ for i in range(num_param_groups):
268
+ merged_partitions = [sd[i] for sd in fp32_flat_groups]
269
+ full_single_fp32_vector = torch.cat(merged_partitions, 0)
270
+ merged_single_partition_of_fp32_groups.append(full_single_fp32_vector)
271
+ avail_numel = sum(
272
+ [full_single_fp32_vector.numel() for full_single_fp32_vector in merged_single_partition_of_fp32_groups])
273
+
274
+ if debug:
275
+ wanted_params = sum([len(shapes) for shapes in param_shapes])
276
+ wanted_numel = sum([sum(shape.numel() for shape in shapes.values()) for shapes in param_shapes])
277
+ # not asserting if there is a mismatch due to possible padding
278
+ print(f"Have {avail_numel} numels to process.")
279
+ print(f"Need {wanted_numel} numels in {wanted_params} params.")
280
+
281
+ # params
282
+ # XXX: for huge models that can't fit into the host's RAM we will have to recode this to support
283
+ # out-of-core computing solution
284
+ total_numel = 0
285
+ total_params = 0
286
+ for shapes, full_single_fp32_vector in zip(param_shapes, merged_single_partition_of_fp32_groups):
287
+ offset = 0
288
+ avail_numel = full_single_fp32_vector.numel()
289
+ for name, shape in shapes.items():
290
+
291
+ unpartitioned_numel = shape.numel() if _has_callable(shape, 'numel') else math.prod(shape)
292
+ total_numel += unpartitioned_numel
293
+ total_params += 1
294
+
295
+ if debug:
296
+ print(f"{name} full shape: {shape} unpartitioned numel {unpartitioned_numel} ")
297
+ state_dict[name] = full_single_fp32_vector.narrow(0, offset, unpartitioned_numel).view(shape)
298
+ offset += unpartitioned_numel
299
+
300
+ # Z2 started to align to 2*world_size to improve nccl performance. Therefore both offset and
301
+ # avail_numel can differ by anywhere between 0..2*world_size. Due to two unrelated complex
302
+ # paddings performed in the code it's almost impossible to predict the exact numbers w/o the
303
+ # live optimizer object, so we are checking that the numbers are within the right range
304
+ align_to = 2 * world_size
305
+
306
+ def zero2_align(x):
307
+ return align_to * math.ceil(x / align_to)
308
+
309
+ if debug:
310
+ print(f"original offset={offset}, avail_numel={avail_numel}")
311
+
312
+ offset = zero2_align(offset)
313
+ avail_numel = zero2_align(avail_numel)
314
+
315
+ if debug:
316
+ print(f"aligned offset={offset}, avail_numel={avail_numel}")
317
+
318
+ # Sanity check
319
+ if offset != avail_numel:
320
+ raise ValueError(f"consumed {offset} numels out of {avail_numel} - something is wrong")
321
+
322
+ print(f"Reconstructed fp32 state dict with {total_params} params {total_numel} elements")
323
+
324
+
325
+ def _get_fp32_state_dict_from_zero2_checkpoint(world_size, fp32_flat_groups, zero_model_states,
326
+ exclude_frozen_parameters):
327
+ state_dict = OrderedDict()
328
+
329
+ # buffers
330
+ buffers = zero_model_states[0].buffers
331
+ state_dict.update(buffers)
332
+ if debug:
333
+ print(f"added {len(buffers)} buffers")
334
+
335
+ if not exclude_frozen_parameters:
336
+ _zero2_merge_frozen_params(state_dict, zero_model_states)
337
+
338
+ _zero2_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states)
339
+
340
+ # recover shared parameters
341
+ for pair in zero_model_states[0].shared_params:
342
+ if pair[1] in state_dict:
343
+ state_dict[pair[0]] = state_dict[pair[1]]
344
+
345
+ return state_dict
346
+
347
+
348
+ def zero3_partitioned_param_info(unpartitioned_numel, world_size):
349
+ remainder = unpartitioned_numel % world_size
350
+ padding_numel = (world_size - remainder) if remainder else 0
351
+ partitioned_numel = math.ceil(unpartitioned_numel / world_size)
352
+ return partitioned_numel, padding_numel
353
+
354
+
355
+ def _zero3_merge_frozen_params(state_dict, world_size, zero_model_states):
356
+ if zero_model_states[0].frozen_param_shapes is None or len(zero_model_states[0].frozen_param_shapes) == 0:
357
+ return
358
+
359
+ if debug:
360
+ for i in range(world_size):
361
+ num_elem = sum(s.numel() for s in zero_model_states[i].frozen_param_fragments.values())
362
+ print(f'rank {i}: {FROZEN_PARAM_SHAPES}.numel = {num_elem}')
363
+
364
+ frozen_param_shapes = zero_model_states[0].frozen_param_shapes
365
+ wanted_params = len(frozen_param_shapes)
366
+ wanted_numel = sum(s.numel() for s in frozen_param_shapes.values())
367
+ avail_numel = sum([p.numel() for p in zero_model_states[0].frozen_param_fragments.values()]) * world_size
368
+ print(f'Frozen params: Have {avail_numel} numels to process.')
369
+ print(f'Frozen params: Need {wanted_numel} numels in {wanted_params} params')
370
+
371
+ total_params = 0
372
+ total_numel = 0
373
+ for name, shape in zero_model_states[0].frozen_param_shapes.items():
374
+ total_params += 1
375
+ unpartitioned_numel = shape.numel()
376
+ total_numel += unpartitioned_numel
377
+
378
+ param_frags = tuple(model_state.frozen_param_fragments[name] for model_state in zero_model_states)
379
+ state_dict[name] = torch.cat(param_frags, 0).narrow(0, 0, unpartitioned_numel).view(shape)
380
+
381
+ partitioned_numel, partitioned_padding_numel = zero3_partitioned_param_info(unpartitioned_numel, world_size)
382
+
383
+ if debug:
384
+ print(
385
+ f"Frozen params: {total_params} {name} full shape: {shape} partition0 numel={partitioned_numel} partitioned_padding_numel={partitioned_padding_numel}"
386
+ )
387
+
388
+ print(f"Reconstructed Frozen fp32 state dict with {total_params} params {total_numel} elements")
389
+
390
+
391
+ class GatheredTensor:
392
+ """
393
+ A pseudo tensor that collects partitioned weights.
394
+ It is more memory efficient when there are multiple groups.
395
+ """
396
+
397
+ def __init__(self, flat_groups, flat_groups_offset, offset, partitioned_numel, shape):
398
+ self.flat_groups = flat_groups
399
+ self.flat_groups_offset = flat_groups_offset
400
+ self.offset = offset
401
+ self.partitioned_numel = partitioned_numel
402
+ self.shape = shape
403
+ self.dtype = self.flat_groups[0][0].dtype
404
+
405
+ def contiguous(self):
406
+ """
407
+ Merge partitioned weights from flat_groups into a single tensor.
408
+ """
409
+ end_idx = self.offset + self.partitioned_numel
410
+ world_size = len(self.flat_groups)
411
+ pad_flat_param_chunks = []
412
+
413
+ for rank_i in range(world_size):
414
+ # for each rank, we need to collect weights from related group/groups
415
+ flat_groups_at_rank_i = self.flat_groups[rank_i]
416
+ start_group_id = None
417
+ end_group_id = None
418
+ for group_id in range(len(self.flat_groups_offset)):
419
+ if self.flat_groups_offset[group_id] <= self.offset < self.flat_groups_offset[group_id + 1]:
420
+ start_group_id = group_id
421
+ if self.flat_groups_offset[group_id] < end_idx <= self.flat_groups_offset[group_id + 1]:
422
+ end_group_id = group_id
423
+ break
424
+ # collect weights from related group/groups
425
+ for group_id in range(start_group_id, end_group_id + 1):
426
+ flat_tensor = flat_groups_at_rank_i[group_id]
427
+ start_offset = self.offset - self.flat_groups_offset[group_id]
428
+ end_offset = min(end_idx, self.flat_groups_offset[group_id + 1]) - self.flat_groups_offset[group_id]
429
+ pad_flat_param_chunks.append(flat_tensor[start_offset:end_offset])
430
+
431
+ # collect weights from all ranks
432
+ pad_flat_param = torch.cat(pad_flat_param_chunks, dim=0)
433
+ param = pad_flat_param[:self.shape.numel()].view(self.shape).contiguous()
434
+ return param
435
+
436
+
437
+ def _zero3_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states):
438
+ param_shapes = zero_model_states[0].param_shapes
439
+ avail_numel = sum([flat_group.numel() for flat_group in fp32_flat_groups[0]]) * world_size
440
+
441
+ # Reconstruction protocol: For zero3 we need to zip the partitions together at boundary of each
442
+ # param, re-consolidating each param, while dealing with padding if any
443
+
444
+ # merge list of dicts, preserving order
445
+ param_shapes = {k: v for d in param_shapes for k, v in d.items()}
446
+
447
+ if debug:
448
+ for i in range(world_size):
449
+ print(f"{FP32_FLAT_GROUPS}[{i}].shape={fp32_flat_groups[i].shape}")
450
+
451
+ wanted_params = len(param_shapes)
452
+ wanted_numel = sum(shape.numel() for shape in param_shapes.values())
453
+ # not asserting if there is a mismatch due to possible padding
454
+ avail_numel = fp32_flat_groups[0].numel() * world_size
455
+ print(f"Trainable params: Have {avail_numel} numels to process.")
456
+ print(f"Trainable params: Need {wanted_numel} numels in {wanted_params} params.")
457
+
458
+ # params
459
+ # XXX: for huge models that can't fit into the host's RAM we will have to recode this to support
460
+ # out-of-core computing solution
461
+ offset = 0
462
+ total_numel = 0
463
+ total_params = 0
464
+ flat_groups_offset = [0] + list(np.cumsum([flat_tensor.numel() for flat_tensor in fp32_flat_groups[0]]))
465
+ for name, shape in tqdm(param_shapes.items(), desc='Gathering sharded weights'):
466
+ unpartitioned_numel = shape.numel()
467
+ total_numel += unpartitioned_numel
468
+ total_params += 1
469
+ partitioned_numel, partitioned_padding_numel = zero3_partitioned_param_info(unpartitioned_numel, world_size)
470
+
471
+ if debug:
472
+ print(
473
+ f"Trainable params: {total_params} {name} full shape: {shape} partition0 numel={partitioned_numel} partitioned_padding_numel={partitioned_padding_numel}"
474
+ )
475
+
476
+ # memory efficient tensor
477
+ tensor = GatheredTensor(fp32_flat_groups, flat_groups_offset, offset, partitioned_numel, shape)
478
+ state_dict[name] = tensor
479
+ offset += partitioned_numel
480
+
481
+ offset *= world_size
482
+
483
+ # Sanity check
484
+ if offset != avail_numel:
485
+ raise ValueError(f"consumed {offset} numels out of {avail_numel} - something is wrong")
486
+
487
+ print(f"Reconstructed Trainable fp32 state dict with {total_params} params {total_numel} elements")
488
+
489
+
490
+ def _get_fp32_state_dict_from_zero3_checkpoint(world_size, fp32_flat_groups, zero_model_states,
491
+ exclude_frozen_parameters):
492
+ state_dict = OrderedDict()
493
+
494
+ # buffers
495
+ buffers = zero_model_states[0].buffers
496
+ state_dict.update(buffers)
497
+ if debug:
498
+ print(f"added {len(buffers)} buffers")
499
+
500
+ if not exclude_frozen_parameters:
501
+ _zero3_merge_frozen_params(state_dict, world_size, zero_model_states)
502
+
503
+ _zero3_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states)
504
+
505
+ # recover shared parameters
506
+ for pair in zero_model_states[0].shared_params:
507
+ if pair[1] in state_dict:
508
+ state_dict[pair[0]] = state_dict[pair[1]]
509
+
510
+ return state_dict
511
+
512
+
513
+ def to_torch_tensor(state_dict, return_empty_tensor=False):
514
+ """
515
+ Convert state_dict of GatheredTensor to torch tensor
516
+ """
517
+ torch_state_dict = {}
518
+ converted_tensors = {}
519
+ for name, tensor in state_dict.items():
520
+ tensor_id = id(tensor)
521
+ if tensor_id in converted_tensors: # shared tensors
522
+ shared_tensor = torch_state_dict[converted_tensors[tensor_id]]
523
+ torch_state_dict[name] = shared_tensor
524
+ else:
525
+ converted_tensors[tensor_id] = name
526
+ if return_empty_tensor:
527
+ torch_state_dict[name] = torch.empty(tensor.shape, dtype=tensor.dtype)
528
+ else:
529
+ torch_state_dict[name] = tensor.contiguous()
530
+ return torch_state_dict
531
+
532
+
533
+ def get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir,
534
+ tag=None,
535
+ exclude_frozen_parameters=False,
536
+ lazy_mode=False):
537
+ """
538
+ Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated state_dict that can be loaded with
539
+ ``load_state_dict()`` and used for training without DeepSpeed or shared with others, for example
540
+ via a model hub.
541
+
542
+ Args:
543
+ - ``checkpoint_dir``: path to the desired checkpoint folder
544
+ - ``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``
545
+ - ``exclude_frozen_parameters``: exclude frozen parameters
546
+ - ``lazy_mode``: get state_dict in lazy mode. It returns a dict of pesduo tensor instead of torch tensor, which is more memory efficient.
547
+ Convert the pesduo tensor to torch tensor by ``.contiguous()``
548
+
549
+ Returns:
550
+ - pytorch ``state_dict``
551
+
552
+ A typical usage might be ::
553
+
554
+ from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint
555
+ # do the training and checkpoint saving
556
+ state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir) # already on cpu
557
+ model = model.cpu() # move to cpu
558
+ model.load_state_dict(state_dict)
559
+ # submit to model hub or save the model to share with others
560
+
561
+ In this example the ``model`` will no longer be usable in the deepspeed context of the same
562
+ application. i.e. you will need to re-initialize the deepspeed engine, since
563
+ ``model.load_state_dict(state_dict)`` will remove all the deepspeed magic from it.
564
+
565
+ If you want it all done for you, use ``load_state_dict_from_zero_checkpoint`` instead.
566
+
567
+ Note: the above usage may not work if your application doesn't have sufficient free CPU memory.
568
+ You may need to use the offline approach using the ``zero_to_fp32.py`` script that is saved with
569
+ the checkpoint. Or you can load state_dict in lazy mode ::
570
+
571
+ from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint
572
+ state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, lazy_mode=True) # not on cpu
573
+ for name, lazy_tensor in state_dict.item():
574
+ tensor = lazy_tensor.contiguous() # to cpu
575
+ print(name, tensor)
576
+ # del tensor to release memory if it no longer in use
577
+ """
578
+ if tag is None:
579
+ latest_path = os.path.join(checkpoint_dir, 'latest')
580
+ if os.path.isfile(latest_path):
581
+ with open(latest_path, 'r') as fd:
582
+ tag = fd.read().strip()
583
+ else:
584
+ raise ValueError(f"Unable to find 'latest' file at {latest_path}")
585
+
586
+ ds_checkpoint_dir = os.path.join(checkpoint_dir, tag)
587
+
588
+ if not os.path.isdir(ds_checkpoint_dir):
589
+ raise FileNotFoundError(f"Directory '{ds_checkpoint_dir}' doesn't exist")
590
+
591
+ state_dict = _get_fp32_state_dict_from_zero_checkpoint(ds_checkpoint_dir, exclude_frozen_parameters)
592
+ if lazy_mode:
593
+ return state_dict
594
+ else:
595
+ return to_torch_tensor(state_dict)
596
+
597
+
598
+ def convert_zero_checkpoint_to_fp32_state_dict(checkpoint_dir,
599
+ output_dir,
600
+ max_shard_size="5GB",
601
+ safe_serialization=False,
602
+ tag=None,
603
+ exclude_frozen_parameters=False):
604
+ """
605
+ Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated ``state_dict`` file that can be
606
+ loaded with ``torch.load(file)`` + ``load_state_dict()`` and used for training without DeepSpeed.
607
+
608
+ Args:
609
+ - ``checkpoint_dir``: path to the desired checkpoint folder. (one that contains the tag-folder, like ``global_step14``)
610
+ - ``output_dir``: directory to the pytorch fp32 state_dict output files
611
+ - ``max_shard_size``: the maximum size for a checkpoint before being sharded, default value is 5GB
612
+ - ``safe_serialization``: whether to save the model using `safetensors` or the traditional PyTorch way (that uses `pickle`).
613
+ - ``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``
614
+ - ``exclude_frozen_parameters``: exclude frozen parameters
615
+ """
616
+
617
+ # Dependency pre-check
618
+ if safe_serialization:
619
+ try:
620
+ from safetensors.torch import save_file
621
+ except ImportError:
622
+ print('If you want to use `safe_serialization`, please `pip install safetensors`')
623
+ raise
624
+ if max_shard_size is not None:
625
+ try:
626
+ from huggingface_hub import split_torch_state_dict_into_shards
627
+ except ImportError:
628
+ print('If you want to use `max_shard_size`, please `pip install huggingface_hub`')
629
+ raise
630
+
631
+ # Convert zero checkpoint to state_dict
632
+ state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir,
633
+ tag,
634
+ exclude_frozen_parameters,
635
+ lazy_mode=True)
636
+
637
+ # Shard the model if it is too big.
638
+ weights_name = "model.safetensors" if safe_serialization else "pytorch_model.bin"
639
+ if max_shard_size is not None:
640
+ filename_pattern = weights_name.replace(".bin", "{suffix}.bin").replace(".safetensors", "{suffix}.safetensors")
641
+ # an memory-efficient approach for sharding
642
+ empty_state_dict = to_torch_tensor(state_dict, return_empty_tensor=True)
643
+ state_dict_split = split_torch_state_dict_into_shards(empty_state_dict,
644
+ filename_pattern=filename_pattern,
645
+ max_shard_size=max_shard_size)
646
+ else:
647
+ from collections import namedtuple
648
+ StateDictSplit = namedtuple("StateDictSplit", ["is_sharded", "filename_to_tensors"])
649
+ state_dict_split = StateDictSplit(is_sharded=False,
650
+ filename_to_tensors={weights_name: list(state_dict.keys())})
651
+
652
+ # Save the model by shard
653
+ os.makedirs(output_dir, exist_ok=True)
654
+ filename_to_tensors = state_dict_split.filename_to_tensors.items()
655
+ for shard_file, tensors in tqdm(filename_to_tensors, desc="Saving checkpoint shards"):
656
+ shard_state_dict = {tensor_name: state_dict[tensor_name] for tensor_name in tensors}
657
+ shard_state_dict = to_torch_tensor(shard_state_dict)
658
+ output_path = os.path.join(output_dir, shard_file)
659
+ if safe_serialization:
660
+ save_file(shard_state_dict, output_path, metadata={"format": "pt"})
661
+ else:
662
+ torch.save(shard_state_dict, output_path)
663
+ # release the memory of current shard
664
+ for tensor_name in list(shard_state_dict.keys()):
665
+ del state_dict[tensor_name]
666
+ del shard_state_dict[tensor_name]
667
+ del shard_state_dict
668
+ gc.collect()
669
+
670
+ # Save index if sharded
671
+ if state_dict_split.is_sharded:
672
+ index = {
673
+ "metadata": state_dict_split.metadata,
674
+ "weight_map": state_dict_split.tensor_to_filename,
675
+ }
676
+ save_index_file = "model.safetensors.index.json" if safe_serialization else "pytorch_model.bin.index.json"
677
+ save_index_file = os.path.join(output_dir, save_index_file)
678
+ with open(save_index_file, "w", encoding="utf-8") as f:
679
+ content = json.dumps(index, indent=2, sort_keys=True) + "\n"
680
+ f.write(content)
681
+
682
+
683
+ def load_state_dict_from_zero_checkpoint(model, checkpoint_dir, tag=None):
684
+ """
685
+ 1. Put the provided model to cpu
686
+ 2. Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated ``state_dict``
687
+ 3. Load it into the provided model
688
+
689
+ Args:
690
+ - ``model``: the model object to update
691
+ - ``checkpoint_dir``: path to the desired checkpoint folder. (one that contains the tag-folder, like ``global_step14``)
692
+ - ``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``
693
+
694
+ Returns:
695
+ - ``model`: modified model
696
+
697
+ Make sure you have plenty of CPU memory available before you call this function. If you don't
698
+ have enough use the ``zero_to_fp32.py`` utility to do the conversion. You will find it
699
+ conveniently placed for you in the checkpoint folder.
700
+
701
+ A typical usage might be ::
702
+
703
+ from deepspeed.utils.zero_to_fp32 import load_state_dict_from_zero_checkpoint
704
+ model = load_state_dict_from_zero_checkpoint(trainer.model, checkpoint_dir)
705
+ # submit to model hub or save the model to share with others
706
+
707
+ Note, that once this was run, the ``model`` will no longer be usable in the deepspeed context
708
+ of the same application. i.e. you will need to re-initialize the deepspeed engine, since
709
+ ``model.load_state_dict(state_dict)`` will remove all the deepspeed magic from it.
710
+
711
+ """
712
+ logger.info("Extracting fp32 weights")
713
+ state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, tag)
714
+
715
+ logger.info("Overwriting model with fp32 weights")
716
+ model = model.cpu()
717
+ model.load_state_dict(state_dict, strict=False)
718
+
719
+ return model
720
+
721
+
722
+ if __name__ == "__main__":
723
+ parser = argparse.ArgumentParser()
724
+ parser.add_argument("checkpoint_dir",
725
+ type=str,
726
+ help="path to the desired checkpoint folder, e.g., path/checkpoint-12")
727
+ parser.add_argument("output_dir",
728
+ type=str,
729
+ help="directory to the pytorch fp32 state_dict output files"
730
+ "(e.g. path/checkpoint-12-output/)")
731
+ parser.add_argument(
732
+ "--max_shard_size",
733
+ type=str,
734
+ default="5GB",
735
+ help="The maximum size for a checkpoint before being sharded. Checkpoints shard will then be each of size"
736
+ "lower than this size. If expressed as a string, needs to be digits followed by a unit (like `5MB`"
737
+ "We default it to 5GB in order for models to be able to run easily on free-tier google colab instances"
738
+ "without CPU OOM issues.")
739
+ parser.add_argument(
740
+ "--safe_serialization",
741
+ default=False,
742
+ action='store_true',
743
+ help="Whether to save the model using `safetensors` or the traditional PyTorch way (that uses `pickle`).")
744
+ parser.add_argument("-t",
745
+ "--tag",
746
+ type=str,
747
+ default=None,
748
+ help="checkpoint tag used as a unique identifier for checkpoint. e.g., global_step1")
749
+ parser.add_argument("--exclude_frozen_parameters", action='store_true', help="exclude frozen parameters")
750
+ parser.add_argument("-d", "--debug", action='store_true', help="enable debug")
751
+ args = parser.parse_args()
752
+
753
+ debug = args.debug
754
+
755
+ convert_zero_checkpoint_to_fp32_state_dict(args.checkpoint_dir,
756
+ args.output_dir,
757
+ max_shard_size=args.max_shard_size,
758
+ safe_serialization=args.safe_serialization,
759
+ tag=args.tag,
760
+ exclude_frozen_parameters=args.exclude_frozen_parameters)
openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-205/config.json ADDED
@@ -0,0 +1,39 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "architectures": [
3
+ "LlamaForCausalLM"
4
+ ],
5
+ "attention_bias": false,
6
+ "attention_dropout": 0.0,
7
+ "bos_token_id": 128000,
8
+ "dtype": "bfloat16",
9
+ "eos_token_id": [
10
+ 128001,
11
+ 128008,
12
+ 128009
13
+ ],
14
+ "head_dim": 64,
15
+ "hidden_act": "silu",
16
+ "hidden_size": 2048,
17
+ "initializer_range": 0.02,
18
+ "intermediate_size": 8192,
19
+ "max_position_embeddings": 131072,
20
+ "mlp_bias": false,
21
+ "model_type": "llama",
22
+ "num_attention_heads": 32,
23
+ "num_hidden_layers": 16,
24
+ "num_key_value_heads": 8,
25
+ "pretraining_tp": 1,
26
+ "rms_norm_eps": 1e-05,
27
+ "rope_scaling": {
28
+ "factor": 32.0,
29
+ "high_freq_factor": 4.0,
30
+ "low_freq_factor": 1.0,
31
+ "original_max_position_embeddings": 8192,
32
+ "rope_type": "llama3"
33
+ },
34
+ "rope_theta": 500000.0,
35
+ "tie_word_embeddings": true,
36
+ "transformers_version": "4.57.1",
37
+ "use_cache": false,
38
+ "vocab_size": 128256
39
+ }
openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-205/generation_config.json ADDED
@@ -0,0 +1,12 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "bos_token_id": 128000,
3
+ "do_sample": true,
4
+ "eos_token_id": [
5
+ 128001,
6
+ 128008,
7
+ 128009
8
+ ],
9
+ "temperature": 0.6,
10
+ "top_p": 0.9,
11
+ "transformers_version": "4.57.1"
12
+ }
openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-205/latest ADDED
@@ -0,0 +1 @@
 
 
1
+ global_step205
openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-205/model.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:fac26a287bd7707ed7900646752bb3ee856f55d81b3d6ef566b39d5fb917d6ca
3
+ size 2996982344
openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-205/scheduler.pt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:9735b7c3044a848443046134ed732a567646f7954e492149d971037508389cf0
3
+ size 1465
openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-205/trainer_state.json ADDED
@@ -0,0 +1,174 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "best_global_step": null,
3
+ "best_metric": null,
4
+ "best_model_checkpoint": null,
5
+ "epoch": 0.2498476538695917,
6
+ "eval_steps": 500,
7
+ "global_step": 205,
8
+ "is_hyper_param_search": false,
9
+ "is_local_process_zero": true,
10
+ "is_world_process_zero": true,
11
+ "log_history": [
12
+ {
13
+ "epoch": 0.01218769043266301,
14
+ "grad_norm": 0.36486390233039856,
15
+ "learning_rate": 5e-05,
16
+ "loss": 1.2213,
17
+ "step": 10
18
+ },
19
+ {
20
+ "epoch": 0.02437538086532602,
21
+ "grad_norm": 0.2953564524650574,
22
+ "learning_rate": 5e-05,
23
+ "loss": 1.0223,
24
+ "step": 20
25
+ },
26
+ {
27
+ "epoch": 0.03656307129798903,
28
+ "grad_norm": 0.2766735851764679,
29
+ "learning_rate": 5e-05,
30
+ "loss": 0.9589,
31
+ "step": 30
32
+ },
33
+ {
34
+ "epoch": 0.04875076173065204,
35
+ "grad_norm": 0.32770997285842896,
36
+ "learning_rate": 5e-05,
37
+ "loss": 0.9081,
38
+ "step": 40
39
+ },
40
+ {
41
+ "epoch": 0.06093845216331505,
42
+ "grad_norm": 0.3052276372909546,
43
+ "learning_rate": 5e-05,
44
+ "loss": 0.8752,
45
+ "step": 50
46
+ },
47
+ {
48
+ "epoch": 0.07312614259597806,
49
+ "grad_norm": 0.240146666765213,
50
+ "learning_rate": 5e-05,
51
+ "loss": 0.873,
52
+ "step": 60
53
+ },
54
+ {
55
+ "epoch": 0.08531383302864107,
56
+ "grad_norm": 0.23872187733650208,
57
+ "learning_rate": 5e-05,
58
+ "loss": 0.8587,
59
+ "step": 70
60
+ },
61
+ {
62
+ "epoch": 0.09750152346130408,
63
+ "grad_norm": 0.32638487219810486,
64
+ "learning_rate": 5e-05,
65
+ "loss": 0.8656,
66
+ "step": 80
67
+ },
68
+ {
69
+ "epoch": 0.10968921389396709,
70
+ "grad_norm": 0.25019028782844543,
71
+ "learning_rate": 5e-05,
72
+ "loss": 0.8685,
73
+ "step": 90
74
+ },
75
+ {
76
+ "epoch": 0.1218769043266301,
77
+ "grad_norm": 0.2507675886154175,
78
+ "learning_rate": 5e-05,
79
+ "loss": 0.8662,
80
+ "step": 100
81
+ },
82
+ {
83
+ "epoch": 0.1340645947592931,
84
+ "grad_norm": 0.27966755628585815,
85
+ "learning_rate": 5e-05,
86
+ "loss": 0.8428,
87
+ "step": 110
88
+ },
89
+ {
90
+ "epoch": 0.14625228519195613,
91
+ "grad_norm": 0.2622673213481903,
92
+ "learning_rate": 5e-05,
93
+ "loss": 0.8416,
94
+ "step": 120
95
+ },
96
+ {
97
+ "epoch": 0.15843997562461914,
98
+ "grad_norm": 0.3003966510295868,
99
+ "learning_rate": 5e-05,
100
+ "loss": 0.8528,
101
+ "step": 130
102
+ },
103
+ {
104
+ "epoch": 0.17062766605728213,
105
+ "grad_norm": 0.32308319211006165,
106
+ "learning_rate": 5e-05,
107
+ "loss": 0.8357,
108
+ "step": 140
109
+ },
110
+ {
111
+ "epoch": 0.18281535648994515,
112
+ "grad_norm": 0.27279096841812134,
113
+ "learning_rate": 5e-05,
114
+ "loss": 0.8432,
115
+ "step": 150
116
+ },
117
+ {
118
+ "epoch": 0.19500304692260817,
119
+ "grad_norm": 0.27976104617118835,
120
+ "learning_rate": 5e-05,
121
+ "loss": 0.8119,
122
+ "step": 160
123
+ },
124
+ {
125
+ "epoch": 0.2071907373552712,
126
+ "grad_norm": 0.290866881608963,
127
+ "learning_rate": 5e-05,
128
+ "loss": 0.8096,
129
+ "step": 170
130
+ },
131
+ {
132
+ "epoch": 0.21937842778793418,
133
+ "grad_norm": 0.29430949687957764,
134
+ "learning_rate": 5e-05,
135
+ "loss": 0.819,
136
+ "step": 180
137
+ },
138
+ {
139
+ "epoch": 0.2315661182205972,
140
+ "grad_norm": 0.27714136242866516,
141
+ "learning_rate": 5e-05,
142
+ "loss": 0.8187,
143
+ "step": 190
144
+ },
145
+ {
146
+ "epoch": 0.2437538086532602,
147
+ "grad_norm": 0.2997029423713684,
148
+ "learning_rate": 5e-05,
149
+ "loss": 0.808,
150
+ "step": 200
151
+ }
152
+ ],
153
+ "logging_steps": 10,
154
+ "max_steps": 2463,
155
+ "num_input_tokens_seen": 0,
156
+ "num_train_epochs": 3,
157
+ "save_steps": 205,
158
+ "stateful_callbacks": {
159
+ "TrainerControl": {
160
+ "args": {
161
+ "should_epoch_stop": false,
162
+ "should_evaluate": false,
163
+ "should_log": false,
164
+ "should_save": true,
165
+ "should_training_stop": false
166
+ },
167
+ "attributes": {}
168
+ }
169
+ },
170
+ "total_flos": 8.38603986725503e+17,
171
+ "train_batch_size": 4,
172
+ "trial_name": null,
173
+ "trial_params": null
174
+ }
openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-205/training_args.bin ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:56896ea21940e33d7e2934267dff2a5b4a4a80195a0f165245f6617a0f88b90b
3
+ size 7377
openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-205/zero_to_fp32.py ADDED
@@ -0,0 +1,760 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python
2
+
3
+ # Copyright (c) Microsoft Corporation.
4
+ # SPDX-License-Identifier: Apache-2.0
5
+
6
+ # DeepSpeed Team
7
+
8
+ # This script extracts fp32 consolidated weights from a zero 1, 2 and 3 DeepSpeed checkpoints. It gets
9
+ # copied into the top level checkpoint dir, so the user can easily do the conversion at any point in
10
+ # the future. Once extracted, the weights don't require DeepSpeed and can be used in any
11
+ # application.
12
+ #
13
+ # example:
14
+ # python zero_to_fp32.py . output_dir/
15
+ # or
16
+ # python zero_to_fp32.py . output_dir/ --safe_serialization
17
+
18
+ import argparse
19
+ import torch
20
+ import glob
21
+ import math
22
+ import os
23
+ import re
24
+ import gc
25
+ import json
26
+ import numpy as np
27
+ from tqdm import tqdm
28
+ from collections import OrderedDict
29
+ from dataclasses import dataclass
30
+
31
+ # while this script doesn't use deepspeed to recover data, since the checkpoints are pickled with
32
+ # DeepSpeed data structures it has to be available in the current python environment.
33
+ from deepspeed.utils import logger
34
+ from deepspeed.checkpoint.constants import (DS_VERSION, OPTIMIZER_STATE_DICT, SINGLE_PARTITION_OF_FP32_GROUPS,
35
+ FP32_FLAT_GROUPS, ZERO_STAGE, PARTITION_COUNT, PARAM_SHAPES, BUFFER_NAMES,
36
+ FROZEN_PARAM_SHAPES, FROZEN_PARAM_FRAGMENTS)
37
+
38
+
39
+ @dataclass
40
+ class zero_model_state:
41
+ buffers: dict()
42
+ param_shapes: dict()
43
+ shared_params: list
44
+ ds_version: int
45
+ frozen_param_shapes: dict()
46
+ frozen_param_fragments: dict()
47
+
48
+
49
+ debug = 0
50
+
51
+ # load to cpu
52
+ device = torch.device('cpu')
53
+
54
+
55
+ def atoi(text):
56
+ return int(text) if text.isdigit() else text
57
+
58
+
59
+ def natural_keys(text):
60
+ '''
61
+ alist.sort(key=natural_keys) sorts in human order
62
+ http://nedbatchelder.com/blog/200712/human_sorting.html
63
+ (See Toothy's implementation in the comments)
64
+ '''
65
+ return [atoi(c) for c in re.split(r'(\d+)', text)]
66
+
67
+
68
+ def get_model_state_file(checkpoint_dir, zero_stage):
69
+ if not os.path.isdir(checkpoint_dir):
70
+ raise FileNotFoundError(f"Directory '{checkpoint_dir}' doesn't exist")
71
+
72
+ # there should be only one file
73
+ if zero_stage <= 2:
74
+ file = os.path.join(checkpoint_dir, "mp_rank_00_model_states.pt")
75
+ elif zero_stage == 3:
76
+ file = os.path.join(checkpoint_dir, "zero_pp_rank_0_mp_rank_00_model_states.pt")
77
+
78
+ if not os.path.exists(file):
79
+ raise FileNotFoundError(f"can't find model states file at '{file}'")
80
+
81
+ return file
82
+
83
+
84
+ def get_checkpoint_files(checkpoint_dir, glob_pattern):
85
+ # XXX: need to test that this simple glob rule works for multi-node setup too
86
+ ckpt_files = sorted(glob.glob(os.path.join(checkpoint_dir, glob_pattern)), key=natural_keys)
87
+
88
+ if len(ckpt_files) == 0:
89
+ raise FileNotFoundError(f"can't find {glob_pattern} files in directory '{checkpoint_dir}'")
90
+
91
+ return ckpt_files
92
+
93
+
94
+ def get_optim_files(checkpoint_dir):
95
+ return get_checkpoint_files(checkpoint_dir, "*_optim_states.pt")
96
+
97
+
98
+ def get_model_state_files(checkpoint_dir):
99
+ return get_checkpoint_files(checkpoint_dir, "*_model_states.pt")
100
+
101
+
102
+ def parse_model_states(files):
103
+ zero_model_states = []
104
+ for file in files:
105
+ state_dict = torch.load(file, map_location=device, weights_only=False)
106
+
107
+ if BUFFER_NAMES not in state_dict:
108
+ raise ValueError(f"{file} is not a model state checkpoint")
109
+ buffer_names = state_dict[BUFFER_NAMES]
110
+ if debug:
111
+ print("Found buffers:", buffer_names)
112
+
113
+ # recover just the buffers while restoring them to fp32 if they were saved in fp16
114
+ buffers = {k: v.float() for k, v in state_dict["module"].items() if k in buffer_names}
115
+ param_shapes = state_dict[PARAM_SHAPES]
116
+
117
+ # collect parameters that are included in param_shapes
118
+ param_names = []
119
+ for s in param_shapes:
120
+ for name in s.keys():
121
+ param_names.append(name)
122
+
123
+ # update with frozen parameters
124
+ frozen_param_shapes = state_dict.get(FROZEN_PARAM_SHAPES, None)
125
+ if frozen_param_shapes is not None:
126
+ if debug:
127
+ print(f"Found frozen_param_shapes: {frozen_param_shapes}")
128
+ param_names += list(frozen_param_shapes.keys())
129
+
130
+ # handle shared params
131
+ shared_params = [[k, v] for k, v in state_dict["shared_params"].items()]
132
+
133
+ ds_version = state_dict.get(DS_VERSION, None)
134
+
135
+ frozen_param_fragments = state_dict.get(FROZEN_PARAM_FRAGMENTS, None)
136
+
137
+ z_model_state = zero_model_state(buffers=buffers,
138
+ param_shapes=param_shapes,
139
+ shared_params=shared_params,
140
+ ds_version=ds_version,
141
+ frozen_param_shapes=frozen_param_shapes,
142
+ frozen_param_fragments=frozen_param_fragments)
143
+ zero_model_states.append(z_model_state)
144
+
145
+ return zero_model_states
146
+
147
+
148
+ def parse_optim_states(files, ds_checkpoint_dir):
149
+ total_files = len(files)
150
+ state_dicts = []
151
+ for f in tqdm(files, desc='Loading checkpoint shards'):
152
+ state_dict = torch.load(f, map_location=device, mmap=True, weights_only=False)
153
+ # immediately discard the potentially huge 2 optimizer states as we only care for fp32 master weights
154
+ # and also handle the case where it was already removed by another helper script
155
+ state_dict["optimizer_state_dict"].pop("optimizer_state_dict", None)
156
+ state_dicts.append(state_dict)
157
+
158
+ if ZERO_STAGE not in state_dicts[0][OPTIMIZER_STATE_DICT]:
159
+ raise ValueError(f"{files[0]} is not a zero checkpoint")
160
+ zero_stage = state_dicts[0][OPTIMIZER_STATE_DICT][ZERO_STAGE]
161
+ world_size = state_dicts[0][OPTIMIZER_STATE_DICT][PARTITION_COUNT]
162
+
163
+ # For ZeRO-2 each param group can have different partition_count as data parallelism for expert
164
+ # parameters can be different from data parallelism for non-expert parameters. So we can just
165
+ # use the max of the partition_count to get the dp world_size.
166
+
167
+ if type(world_size) is list:
168
+ world_size = max(world_size)
169
+
170
+ if world_size != total_files:
171
+ raise ValueError(
172
+ f"Expected {world_size} of '*_optim_states.pt' under '{ds_checkpoint_dir}' but found {total_files} files. "
173
+ "Possibly due to an overwrite of an old checkpoint, or a checkpoint didn't get saved by one or more processes."
174
+ )
175
+
176
+ # the groups are named differently in each stage
177
+ if zero_stage <= 2:
178
+ fp32_groups_key = SINGLE_PARTITION_OF_FP32_GROUPS
179
+ elif zero_stage == 3:
180
+ fp32_groups_key = FP32_FLAT_GROUPS
181
+ else:
182
+ raise ValueError(f"unknown zero stage {zero_stage}")
183
+
184
+ fp32_flat_groups = [state_dicts[i][OPTIMIZER_STATE_DICT][fp32_groups_key] for i in range(len(state_dicts))]
185
+ return zero_stage, world_size, fp32_flat_groups
186
+
187
+
188
+ def _get_fp32_state_dict_from_zero_checkpoint(ds_checkpoint_dir, exclude_frozen_parameters):
189
+ """
190
+ Returns fp32 state_dict reconstructed from ds checkpoint
191
+
192
+ Args:
193
+ - ``ds_checkpoint_dir``: path to the deepspeed checkpoint folder (where the optimizer files are)
194
+
195
+ """
196
+ print(f"Processing zero checkpoint '{ds_checkpoint_dir}'")
197
+
198
+ optim_files = get_optim_files(ds_checkpoint_dir)
199
+ zero_stage, world_size, fp32_flat_groups = parse_optim_states(optim_files, ds_checkpoint_dir)
200
+ print(f"Detected checkpoint of type zero stage {zero_stage}, world_size: {world_size}")
201
+
202
+ model_files = get_model_state_files(ds_checkpoint_dir)
203
+
204
+ zero_model_states = parse_model_states(model_files)
205
+ print(f'Parsing checkpoint created by deepspeed=={zero_model_states[0].ds_version}')
206
+
207
+ if zero_stage <= 2:
208
+ return _get_fp32_state_dict_from_zero2_checkpoint(world_size, fp32_flat_groups, zero_model_states,
209
+ exclude_frozen_parameters)
210
+ elif zero_stage == 3:
211
+ return _get_fp32_state_dict_from_zero3_checkpoint(world_size, fp32_flat_groups, zero_model_states,
212
+ exclude_frozen_parameters)
213
+
214
+
215
+ def _zero2_merge_frozen_params(state_dict, zero_model_states):
216
+ if zero_model_states[0].frozen_param_shapes is None or len(zero_model_states[0].frozen_param_shapes) == 0:
217
+ return
218
+
219
+ frozen_param_shapes = zero_model_states[0].frozen_param_shapes
220
+ frozen_param_fragments = zero_model_states[0].frozen_param_fragments
221
+
222
+ if debug:
223
+ num_elem = sum(s.numel() for s in frozen_param_shapes.values())
224
+ print(f'rank 0: {FROZEN_PARAM_SHAPES}.numel = {num_elem}')
225
+
226
+ wanted_params = len(frozen_param_shapes)
227
+ wanted_numel = sum(s.numel() for s in frozen_param_shapes.values())
228
+ avail_numel = sum([p.numel() for p in frozen_param_fragments.values()])
229
+ print(f'Frozen params: Have {avail_numel} numels to process.')
230
+ print(f'Frozen params: Need {wanted_numel} numels in {wanted_params} params')
231
+
232
+ total_params = 0
233
+ total_numel = 0
234
+ for name, shape in frozen_param_shapes.items():
235
+ total_params += 1
236
+ unpartitioned_numel = shape.numel()
237
+ total_numel += unpartitioned_numel
238
+
239
+ state_dict[name] = frozen_param_fragments[name]
240
+
241
+ if debug:
242
+ print(f"{name} full shape: {shape} unpartitioned numel {unpartitioned_numel} ")
243
+
244
+ print(f"Reconstructed Frozen fp32 state dict with {total_params} params {total_numel} elements")
245
+
246
+
247
+ def _has_callable(obj, fn):
248
+ attr = getattr(obj, fn, None)
249
+ return callable(attr)
250
+
251
+
252
+ def _zero2_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states):
253
+ param_shapes = zero_model_states[0].param_shapes
254
+
255
+ # Reconstruction protocol:
256
+ #
257
+ # XXX: document this
258
+
259
+ if debug:
260
+ for i in range(world_size):
261
+ for j in range(len(fp32_flat_groups[0])):
262
+ print(f"{FP32_FLAT_GROUPS}[{i}][{j}].shape={fp32_flat_groups[i][j].shape}")
263
+
264
+ # XXX: memory usage doubles here (zero2)
265
+ num_param_groups = len(fp32_flat_groups[0])
266
+ merged_single_partition_of_fp32_groups = []
267
+ for i in range(num_param_groups):
268
+ merged_partitions = [sd[i] for sd in fp32_flat_groups]
269
+ full_single_fp32_vector = torch.cat(merged_partitions, 0)
270
+ merged_single_partition_of_fp32_groups.append(full_single_fp32_vector)
271
+ avail_numel = sum(
272
+ [full_single_fp32_vector.numel() for full_single_fp32_vector in merged_single_partition_of_fp32_groups])
273
+
274
+ if debug:
275
+ wanted_params = sum([len(shapes) for shapes in param_shapes])
276
+ wanted_numel = sum([sum(shape.numel() for shape in shapes.values()) for shapes in param_shapes])
277
+ # not asserting if there is a mismatch due to possible padding
278
+ print(f"Have {avail_numel} numels to process.")
279
+ print(f"Need {wanted_numel} numels in {wanted_params} params.")
280
+
281
+ # params
282
+ # XXX: for huge models that can't fit into the host's RAM we will have to recode this to support
283
+ # out-of-core computing solution
284
+ total_numel = 0
285
+ total_params = 0
286
+ for shapes, full_single_fp32_vector in zip(param_shapes, merged_single_partition_of_fp32_groups):
287
+ offset = 0
288
+ avail_numel = full_single_fp32_vector.numel()
289
+ for name, shape in shapes.items():
290
+
291
+ unpartitioned_numel = shape.numel() if _has_callable(shape, 'numel') else math.prod(shape)
292
+ total_numel += unpartitioned_numel
293
+ total_params += 1
294
+
295
+ if debug:
296
+ print(f"{name} full shape: {shape} unpartitioned numel {unpartitioned_numel} ")
297
+ state_dict[name] = full_single_fp32_vector.narrow(0, offset, unpartitioned_numel).view(shape)
298
+ offset += unpartitioned_numel
299
+
300
+ # Z2 started to align to 2*world_size to improve nccl performance. Therefore both offset and
301
+ # avail_numel can differ by anywhere between 0..2*world_size. Due to two unrelated complex
302
+ # paddings performed in the code it's almost impossible to predict the exact numbers w/o the
303
+ # live optimizer object, so we are checking that the numbers are within the right range
304
+ align_to = 2 * world_size
305
+
306
+ def zero2_align(x):
307
+ return align_to * math.ceil(x / align_to)
308
+
309
+ if debug:
310
+ print(f"original offset={offset}, avail_numel={avail_numel}")
311
+
312
+ offset = zero2_align(offset)
313
+ avail_numel = zero2_align(avail_numel)
314
+
315
+ if debug:
316
+ print(f"aligned offset={offset}, avail_numel={avail_numel}")
317
+
318
+ # Sanity check
319
+ if offset != avail_numel:
320
+ raise ValueError(f"consumed {offset} numels out of {avail_numel} - something is wrong")
321
+
322
+ print(f"Reconstructed fp32 state dict with {total_params} params {total_numel} elements")
323
+
324
+
325
+ def _get_fp32_state_dict_from_zero2_checkpoint(world_size, fp32_flat_groups, zero_model_states,
326
+ exclude_frozen_parameters):
327
+ state_dict = OrderedDict()
328
+
329
+ # buffers
330
+ buffers = zero_model_states[0].buffers
331
+ state_dict.update(buffers)
332
+ if debug:
333
+ print(f"added {len(buffers)} buffers")
334
+
335
+ if not exclude_frozen_parameters:
336
+ _zero2_merge_frozen_params(state_dict, zero_model_states)
337
+
338
+ _zero2_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states)
339
+
340
+ # recover shared parameters
341
+ for pair in zero_model_states[0].shared_params:
342
+ if pair[1] in state_dict:
343
+ state_dict[pair[0]] = state_dict[pair[1]]
344
+
345
+ return state_dict
346
+
347
+
348
+ def zero3_partitioned_param_info(unpartitioned_numel, world_size):
349
+ remainder = unpartitioned_numel % world_size
350
+ padding_numel = (world_size - remainder) if remainder else 0
351
+ partitioned_numel = math.ceil(unpartitioned_numel / world_size)
352
+ return partitioned_numel, padding_numel
353
+
354
+
355
+ def _zero3_merge_frozen_params(state_dict, world_size, zero_model_states):
356
+ if zero_model_states[0].frozen_param_shapes is None or len(zero_model_states[0].frozen_param_shapes) == 0:
357
+ return
358
+
359
+ if debug:
360
+ for i in range(world_size):
361
+ num_elem = sum(s.numel() for s in zero_model_states[i].frozen_param_fragments.values())
362
+ print(f'rank {i}: {FROZEN_PARAM_SHAPES}.numel = {num_elem}')
363
+
364
+ frozen_param_shapes = zero_model_states[0].frozen_param_shapes
365
+ wanted_params = len(frozen_param_shapes)
366
+ wanted_numel = sum(s.numel() for s in frozen_param_shapes.values())
367
+ avail_numel = sum([p.numel() for p in zero_model_states[0].frozen_param_fragments.values()]) * world_size
368
+ print(f'Frozen params: Have {avail_numel} numels to process.')
369
+ print(f'Frozen params: Need {wanted_numel} numels in {wanted_params} params')
370
+
371
+ total_params = 0
372
+ total_numel = 0
373
+ for name, shape in zero_model_states[0].frozen_param_shapes.items():
374
+ total_params += 1
375
+ unpartitioned_numel = shape.numel()
376
+ total_numel += unpartitioned_numel
377
+
378
+ param_frags = tuple(model_state.frozen_param_fragments[name] for model_state in zero_model_states)
379
+ state_dict[name] = torch.cat(param_frags, 0).narrow(0, 0, unpartitioned_numel).view(shape)
380
+
381
+ partitioned_numel, partitioned_padding_numel = zero3_partitioned_param_info(unpartitioned_numel, world_size)
382
+
383
+ if debug:
384
+ print(
385
+ f"Frozen params: {total_params} {name} full shape: {shape} partition0 numel={partitioned_numel} partitioned_padding_numel={partitioned_padding_numel}"
386
+ )
387
+
388
+ print(f"Reconstructed Frozen fp32 state dict with {total_params} params {total_numel} elements")
389
+
390
+
391
+ class GatheredTensor:
392
+ """
393
+ A pseudo tensor that collects partitioned weights.
394
+ It is more memory efficient when there are multiple groups.
395
+ """
396
+
397
+ def __init__(self, flat_groups, flat_groups_offset, offset, partitioned_numel, shape):
398
+ self.flat_groups = flat_groups
399
+ self.flat_groups_offset = flat_groups_offset
400
+ self.offset = offset
401
+ self.partitioned_numel = partitioned_numel
402
+ self.shape = shape
403
+ self.dtype = self.flat_groups[0][0].dtype
404
+
405
+ def contiguous(self):
406
+ """
407
+ Merge partitioned weights from flat_groups into a single tensor.
408
+ """
409
+ end_idx = self.offset + self.partitioned_numel
410
+ world_size = len(self.flat_groups)
411
+ pad_flat_param_chunks = []
412
+
413
+ for rank_i in range(world_size):
414
+ # for each rank, we need to collect weights from related group/groups
415
+ flat_groups_at_rank_i = self.flat_groups[rank_i]
416
+ start_group_id = None
417
+ end_group_id = None
418
+ for group_id in range(len(self.flat_groups_offset)):
419
+ if self.flat_groups_offset[group_id] <= self.offset < self.flat_groups_offset[group_id + 1]:
420
+ start_group_id = group_id
421
+ if self.flat_groups_offset[group_id] < end_idx <= self.flat_groups_offset[group_id + 1]:
422
+ end_group_id = group_id
423
+ break
424
+ # collect weights from related group/groups
425
+ for group_id in range(start_group_id, end_group_id + 1):
426
+ flat_tensor = flat_groups_at_rank_i[group_id]
427
+ start_offset = self.offset - self.flat_groups_offset[group_id]
428
+ end_offset = min(end_idx, self.flat_groups_offset[group_id + 1]) - self.flat_groups_offset[group_id]
429
+ pad_flat_param_chunks.append(flat_tensor[start_offset:end_offset])
430
+
431
+ # collect weights from all ranks
432
+ pad_flat_param = torch.cat(pad_flat_param_chunks, dim=0)
433
+ param = pad_flat_param[:self.shape.numel()].view(self.shape).contiguous()
434
+ return param
435
+
436
+
437
+ def _zero3_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states):
438
+ param_shapes = zero_model_states[0].param_shapes
439
+ avail_numel = sum([flat_group.numel() for flat_group in fp32_flat_groups[0]]) * world_size
440
+
441
+ # Reconstruction protocol: For zero3 we need to zip the partitions together at boundary of each
442
+ # param, re-consolidating each param, while dealing with padding if any
443
+
444
+ # merge list of dicts, preserving order
445
+ param_shapes = {k: v for d in param_shapes for k, v in d.items()}
446
+
447
+ if debug:
448
+ for i in range(world_size):
449
+ print(f"{FP32_FLAT_GROUPS}[{i}].shape={fp32_flat_groups[i].shape}")
450
+
451
+ wanted_params = len(param_shapes)
452
+ wanted_numel = sum(shape.numel() for shape in param_shapes.values())
453
+ # not asserting if there is a mismatch due to possible padding
454
+ avail_numel = fp32_flat_groups[0].numel() * world_size
455
+ print(f"Trainable params: Have {avail_numel} numels to process.")
456
+ print(f"Trainable params: Need {wanted_numel} numels in {wanted_params} params.")
457
+
458
+ # params
459
+ # XXX: for huge models that can't fit into the host's RAM we will have to recode this to support
460
+ # out-of-core computing solution
461
+ offset = 0
462
+ total_numel = 0
463
+ total_params = 0
464
+ flat_groups_offset = [0] + list(np.cumsum([flat_tensor.numel() for flat_tensor in fp32_flat_groups[0]]))
465
+ for name, shape in tqdm(param_shapes.items(), desc='Gathering sharded weights'):
466
+ unpartitioned_numel = shape.numel()
467
+ total_numel += unpartitioned_numel
468
+ total_params += 1
469
+ partitioned_numel, partitioned_padding_numel = zero3_partitioned_param_info(unpartitioned_numel, world_size)
470
+
471
+ if debug:
472
+ print(
473
+ f"Trainable params: {total_params} {name} full shape: {shape} partition0 numel={partitioned_numel} partitioned_padding_numel={partitioned_padding_numel}"
474
+ )
475
+
476
+ # memory efficient tensor
477
+ tensor = GatheredTensor(fp32_flat_groups, flat_groups_offset, offset, partitioned_numel, shape)
478
+ state_dict[name] = tensor
479
+ offset += partitioned_numel
480
+
481
+ offset *= world_size
482
+
483
+ # Sanity check
484
+ if offset != avail_numel:
485
+ raise ValueError(f"consumed {offset} numels out of {avail_numel} - something is wrong")
486
+
487
+ print(f"Reconstructed Trainable fp32 state dict with {total_params} params {total_numel} elements")
488
+
489
+
490
+ def _get_fp32_state_dict_from_zero3_checkpoint(world_size, fp32_flat_groups, zero_model_states,
491
+ exclude_frozen_parameters):
492
+ state_dict = OrderedDict()
493
+
494
+ # buffers
495
+ buffers = zero_model_states[0].buffers
496
+ state_dict.update(buffers)
497
+ if debug:
498
+ print(f"added {len(buffers)} buffers")
499
+
500
+ if not exclude_frozen_parameters:
501
+ _zero3_merge_frozen_params(state_dict, world_size, zero_model_states)
502
+
503
+ _zero3_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states)
504
+
505
+ # recover shared parameters
506
+ for pair in zero_model_states[0].shared_params:
507
+ if pair[1] in state_dict:
508
+ state_dict[pair[0]] = state_dict[pair[1]]
509
+
510
+ return state_dict
511
+
512
+
513
+ def to_torch_tensor(state_dict, return_empty_tensor=False):
514
+ """
515
+ Convert state_dict of GatheredTensor to torch tensor
516
+ """
517
+ torch_state_dict = {}
518
+ converted_tensors = {}
519
+ for name, tensor in state_dict.items():
520
+ tensor_id = id(tensor)
521
+ if tensor_id in converted_tensors: # shared tensors
522
+ shared_tensor = torch_state_dict[converted_tensors[tensor_id]]
523
+ torch_state_dict[name] = shared_tensor
524
+ else:
525
+ converted_tensors[tensor_id] = name
526
+ if return_empty_tensor:
527
+ torch_state_dict[name] = torch.empty(tensor.shape, dtype=tensor.dtype)
528
+ else:
529
+ torch_state_dict[name] = tensor.contiguous()
530
+ return torch_state_dict
531
+
532
+
533
+ def get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir,
534
+ tag=None,
535
+ exclude_frozen_parameters=False,
536
+ lazy_mode=False):
537
+ """
538
+ Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated state_dict that can be loaded with
539
+ ``load_state_dict()`` and used for training without DeepSpeed or shared with others, for example
540
+ via a model hub.
541
+
542
+ Args:
543
+ - ``checkpoint_dir``: path to the desired checkpoint folder
544
+ - ``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``
545
+ - ``exclude_frozen_parameters``: exclude frozen parameters
546
+ - ``lazy_mode``: get state_dict in lazy mode. It returns a dict of pesduo tensor instead of torch tensor, which is more memory efficient.
547
+ Convert the pesduo tensor to torch tensor by ``.contiguous()``
548
+
549
+ Returns:
550
+ - pytorch ``state_dict``
551
+
552
+ A typical usage might be ::
553
+
554
+ from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint
555
+ # do the training and checkpoint saving
556
+ state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir) # already on cpu
557
+ model = model.cpu() # move to cpu
558
+ model.load_state_dict(state_dict)
559
+ # submit to model hub or save the model to share with others
560
+
561
+ In this example the ``model`` will no longer be usable in the deepspeed context of the same
562
+ application. i.e. you will need to re-initialize the deepspeed engine, since
563
+ ``model.load_state_dict(state_dict)`` will remove all the deepspeed magic from it.
564
+
565
+ If you want it all done for you, use ``load_state_dict_from_zero_checkpoint`` instead.
566
+
567
+ Note: the above usage may not work if your application doesn't have sufficient free CPU memory.
568
+ You may need to use the offline approach using the ``zero_to_fp32.py`` script that is saved with
569
+ the checkpoint. Or you can load state_dict in lazy mode ::
570
+
571
+ from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint
572
+ state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, lazy_mode=True) # not on cpu
573
+ for name, lazy_tensor in state_dict.item():
574
+ tensor = lazy_tensor.contiguous() # to cpu
575
+ print(name, tensor)
576
+ # del tensor to release memory if it no longer in use
577
+ """
578
+ if tag is None:
579
+ latest_path = os.path.join(checkpoint_dir, 'latest')
580
+ if os.path.isfile(latest_path):
581
+ with open(latest_path, 'r') as fd:
582
+ tag = fd.read().strip()
583
+ else:
584
+ raise ValueError(f"Unable to find 'latest' file at {latest_path}")
585
+
586
+ ds_checkpoint_dir = os.path.join(checkpoint_dir, tag)
587
+
588
+ if not os.path.isdir(ds_checkpoint_dir):
589
+ raise FileNotFoundError(f"Directory '{ds_checkpoint_dir}' doesn't exist")
590
+
591
+ state_dict = _get_fp32_state_dict_from_zero_checkpoint(ds_checkpoint_dir, exclude_frozen_parameters)
592
+ if lazy_mode:
593
+ return state_dict
594
+ else:
595
+ return to_torch_tensor(state_dict)
596
+
597
+
598
+ def convert_zero_checkpoint_to_fp32_state_dict(checkpoint_dir,
599
+ output_dir,
600
+ max_shard_size="5GB",
601
+ safe_serialization=False,
602
+ tag=None,
603
+ exclude_frozen_parameters=False):
604
+ """
605
+ Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated ``state_dict`` file that can be
606
+ loaded with ``torch.load(file)`` + ``load_state_dict()`` and used for training without DeepSpeed.
607
+
608
+ Args:
609
+ - ``checkpoint_dir``: path to the desired checkpoint folder. (one that contains the tag-folder, like ``global_step14``)
610
+ - ``output_dir``: directory to the pytorch fp32 state_dict output files
611
+ - ``max_shard_size``: the maximum size for a checkpoint before being sharded, default value is 5GB
612
+ - ``safe_serialization``: whether to save the model using `safetensors` or the traditional PyTorch way (that uses `pickle`).
613
+ - ``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``
614
+ - ``exclude_frozen_parameters``: exclude frozen parameters
615
+ """
616
+
617
+ # Dependency pre-check
618
+ if safe_serialization:
619
+ try:
620
+ from safetensors.torch import save_file
621
+ except ImportError:
622
+ print('If you want to use `safe_serialization`, please `pip install safetensors`')
623
+ raise
624
+ if max_shard_size is not None:
625
+ try:
626
+ from huggingface_hub import split_torch_state_dict_into_shards
627
+ except ImportError:
628
+ print('If you want to use `max_shard_size`, please `pip install huggingface_hub`')
629
+ raise
630
+
631
+ # Convert zero checkpoint to state_dict
632
+ state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir,
633
+ tag,
634
+ exclude_frozen_parameters,
635
+ lazy_mode=True)
636
+
637
+ # Shard the model if it is too big.
638
+ weights_name = "model.safetensors" if safe_serialization else "pytorch_model.bin"
639
+ if max_shard_size is not None:
640
+ filename_pattern = weights_name.replace(".bin", "{suffix}.bin").replace(".safetensors", "{suffix}.safetensors")
641
+ # an memory-efficient approach for sharding
642
+ empty_state_dict = to_torch_tensor(state_dict, return_empty_tensor=True)
643
+ state_dict_split = split_torch_state_dict_into_shards(empty_state_dict,
644
+ filename_pattern=filename_pattern,
645
+ max_shard_size=max_shard_size)
646
+ else:
647
+ from collections import namedtuple
648
+ StateDictSplit = namedtuple("StateDictSplit", ["is_sharded", "filename_to_tensors"])
649
+ state_dict_split = StateDictSplit(is_sharded=False,
650
+ filename_to_tensors={weights_name: list(state_dict.keys())})
651
+
652
+ # Save the model by shard
653
+ os.makedirs(output_dir, exist_ok=True)
654
+ filename_to_tensors = state_dict_split.filename_to_tensors.items()
655
+ for shard_file, tensors in tqdm(filename_to_tensors, desc="Saving checkpoint shards"):
656
+ shard_state_dict = {tensor_name: state_dict[tensor_name] for tensor_name in tensors}
657
+ shard_state_dict = to_torch_tensor(shard_state_dict)
658
+ output_path = os.path.join(output_dir, shard_file)
659
+ if safe_serialization:
660
+ save_file(shard_state_dict, output_path, metadata={"format": "pt"})
661
+ else:
662
+ torch.save(shard_state_dict, output_path)
663
+ # release the memory of current shard
664
+ for tensor_name in list(shard_state_dict.keys()):
665
+ del state_dict[tensor_name]
666
+ del shard_state_dict[tensor_name]
667
+ del shard_state_dict
668
+ gc.collect()
669
+
670
+ # Save index if sharded
671
+ if state_dict_split.is_sharded:
672
+ index = {
673
+ "metadata": state_dict_split.metadata,
674
+ "weight_map": state_dict_split.tensor_to_filename,
675
+ }
676
+ save_index_file = "model.safetensors.index.json" if safe_serialization else "pytorch_model.bin.index.json"
677
+ save_index_file = os.path.join(output_dir, save_index_file)
678
+ with open(save_index_file, "w", encoding="utf-8") as f:
679
+ content = json.dumps(index, indent=2, sort_keys=True) + "\n"
680
+ f.write(content)
681
+
682
+
683
+ def load_state_dict_from_zero_checkpoint(model, checkpoint_dir, tag=None):
684
+ """
685
+ 1. Put the provided model to cpu
686
+ 2. Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated ``state_dict``
687
+ 3. Load it into the provided model
688
+
689
+ Args:
690
+ - ``model``: the model object to update
691
+ - ``checkpoint_dir``: path to the desired checkpoint folder. (one that contains the tag-folder, like ``global_step14``)
692
+ - ``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``
693
+
694
+ Returns:
695
+ - ``model`: modified model
696
+
697
+ Make sure you have plenty of CPU memory available before you call this function. If you don't
698
+ have enough use the ``zero_to_fp32.py`` utility to do the conversion. You will find it
699
+ conveniently placed for you in the checkpoint folder.
700
+
701
+ A typical usage might be ::
702
+
703
+ from deepspeed.utils.zero_to_fp32 import load_state_dict_from_zero_checkpoint
704
+ model = load_state_dict_from_zero_checkpoint(trainer.model, checkpoint_dir)
705
+ # submit to model hub or save the model to share with others
706
+
707
+ Note, that once this was run, the ``model`` will no longer be usable in the deepspeed context
708
+ of the same application. i.e. you will need to re-initialize the deepspeed engine, since
709
+ ``model.load_state_dict(state_dict)`` will remove all the deepspeed magic from it.
710
+
711
+ """
712
+ logger.info("Extracting fp32 weights")
713
+ state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, tag)
714
+
715
+ logger.info("Overwriting model with fp32 weights")
716
+ model = model.cpu()
717
+ model.load_state_dict(state_dict, strict=False)
718
+
719
+ return model
720
+
721
+
722
+ if __name__ == "__main__":
723
+ parser = argparse.ArgumentParser()
724
+ parser.add_argument("checkpoint_dir",
725
+ type=str,
726
+ help="path to the desired checkpoint folder, e.g., path/checkpoint-12")
727
+ parser.add_argument("output_dir",
728
+ type=str,
729
+ help="directory to the pytorch fp32 state_dict output files"
730
+ "(e.g. path/checkpoint-12-output/)")
731
+ parser.add_argument(
732
+ "--max_shard_size",
733
+ type=str,
734
+ default="5GB",
735
+ help="The maximum size for a checkpoint before being sharded. Checkpoints shard will then be each of size"
736
+ "lower than this size. If expressed as a string, needs to be digits followed by a unit (like `5MB`"
737
+ "We default it to 5GB in order for models to be able to run easily on free-tier google colab instances"
738
+ "without CPU OOM issues.")
739
+ parser.add_argument(
740
+ "--safe_serialization",
741
+ default=False,
742
+ action='store_true',
743
+ help="Whether to save the model using `safetensors` or the traditional PyTorch way (that uses `pickle`).")
744
+ parser.add_argument("-t",
745
+ "--tag",
746
+ type=str,
747
+ default=None,
748
+ help="checkpoint tag used as a unique identifier for checkpoint. e.g., global_step1")
749
+ parser.add_argument("--exclude_frozen_parameters", action='store_true', help="exclude frozen parameters")
750
+ parser.add_argument("-d", "--debug", action='store_true', help="enable debug")
751
+ args = parser.parse_args()
752
+
753
+ debug = args.debug
754
+
755
+ convert_zero_checkpoint_to_fp32_state_dict(args.checkpoint_dir,
756
+ args.output_dir,
757
+ max_shard_size=args.max_shard_size,
758
+ safe_serialization=args.safe_serialization,
759
+ tag=args.tag,
760
+ exclude_frozen_parameters=args.exclude_frozen_parameters)
openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-2050/config.json ADDED
@@ -0,0 +1,39 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "architectures": [
3
+ "LlamaForCausalLM"
4
+ ],
5
+ "attention_bias": false,
6
+ "attention_dropout": 0.0,
7
+ "bos_token_id": 128000,
8
+ "dtype": "bfloat16",
9
+ "eos_token_id": [
10
+ 128001,
11
+ 128008,
12
+ 128009
13
+ ],
14
+ "head_dim": 64,
15
+ "hidden_act": "silu",
16
+ "hidden_size": 2048,
17
+ "initializer_range": 0.02,
18
+ "intermediate_size": 8192,
19
+ "max_position_embeddings": 131072,
20
+ "mlp_bias": false,
21
+ "model_type": "llama",
22
+ "num_attention_heads": 32,
23
+ "num_hidden_layers": 16,
24
+ "num_key_value_heads": 8,
25
+ "pretraining_tp": 1,
26
+ "rms_norm_eps": 1e-05,
27
+ "rope_scaling": {
28
+ "factor": 32.0,
29
+ "high_freq_factor": 4.0,
30
+ "low_freq_factor": 1.0,
31
+ "original_max_position_embeddings": 8192,
32
+ "rope_type": "llama3"
33
+ },
34
+ "rope_theta": 500000.0,
35
+ "tie_word_embeddings": true,
36
+ "transformers_version": "4.57.1",
37
+ "use_cache": false,
38
+ "vocab_size": 128256
39
+ }
openai_gpt_oss_120b_Qwen_Qwen3_235B_A22B_fp8_tput_deepseek_ai_DeepSeek_R1_Distill_Qwen_14B/checkpoint-2050/generation_config.json ADDED
@@ -0,0 +1,12 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "bos_token_id": 128000,
3
+ "do_sample": true,
4
+ "eos_token_id": [
5
+ 128001,
6
+ 128008,
7
+ 128009
8
+ ],
9
+ "temperature": 0.6,
10
+ "top_p": 0.9,
11
+ "transformers_version": "4.57.1"
12
+ }