Delete else

else/added_tokens.json

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

else/trainer_state.json

@@ -1,1833 +0,0 @@
-{
-  "best_metric": null,
-  "best_model_checkpoint": null,
-  "epoch": 0.02529724260055654,
-  "eval_steps": 500,
-  "global_step": 100,
-  "is_hyper_param_search": false,
-  "is_local_process_zero": true,
-  "is_world_process_zero": true,
-  "log_history": [
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 440.0,
-      "epoch": 0.00025297242600556537,
-      "grad_norm": 99.50328063964844,
-      "kl": 0.0003566741943359375,
-      "learning_rate": 9.99873513786997e-07,
-      "loss": 0.0,
-      "reward": 5.4594926834106445,
-      "reward_std": 0.18958747386932373,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.7995114922523499,
-      "rewards/location_misalighment_heatmap_reward": 0.8512309193611145,
-      "rewards/overall_score_rewards": 0.8087499141693115,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 1
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 496.75,
-      "epoch": 0.0005059448520111307,
-      "grad_norm": 2.2404987812042236,
-      "kl": 0.0003204345703125,
-      "learning_rate": 9.997470275739944e-07,
-      "loss": 0.0,
-      "reward": 5.263876438140869,
-      "reward_std": 0.18220201134681702,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.5646694898605347,
-      "rewards/location_misalighment_heatmap_reward": 0.822956919670105,
-      "rewards/overall_score_rewards": 0.8762500286102295,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 2
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 512.5,
-      "epoch": 0.0007589172780166962,
-      "grad_norm": 3.8491673469543457,
-      "kl": 0.000591278076171875,
-      "learning_rate": 9.996205413609917e-07,
-      "loss": -0.0003,
-      "reward": 5.384895324707031,
-      "reward_std": 0.12690222263336182,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.7779643535614014,
-      "rewards/location_misalighment_heatmap_reward": 0.7681811451911926,
-      "rewards/overall_score_rewards": 0.8387500047683716,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 3
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 460.25,
-      "epoch": 0.0010118897040222615,
-      "grad_norm": 3.278803825378418,
-      "kl": 0.00078582763671875,
-      "learning_rate": 9.994940551479888e-07,
-      "loss": 0.0,
-      "reward": 5.464195251464844,
-      "reward_std": 0.17130990326404572,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.838955283164978,
-      "rewards/location_misalighment_heatmap_reward": 0.838990330696106,
-      "rewards/overall_score_rewards": 0.7862499952316284,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 4
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 516.875,
-      "epoch": 0.0012648621300278269,
-      "grad_norm": 3.218126058578491,
-      "kl": 0.000507354736328125,
-      "learning_rate": 9.99367568934986e-07,
-      "loss": 0.0,
-      "reward": 5.111246109008789,
-      "reward_std": 0.05334939435124397,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.5527762770652771,
-      "rewards/location_misalighment_heatmap_reward": 0.7759696245193481,
-      "rewards/overall_score_rewards": 0.7825000286102295,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 5
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 495.125,
-      "epoch": 0.0015178345560333923,
-      "grad_norm": 3.257352352142334,
-      "kl": 0.0004634857177734375,
-      "learning_rate": 9.992410827219832e-07,
-      "loss": 0.0,
-      "reward": 5.156436920166016,
-      "reward_std": 0.1203070878982544,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.46342724561691284,
-      "rewards/location_misalighment_heatmap_reward": 0.850509524345398,
-      "rewards/overall_score_rewards": 0.8425000905990601,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 6
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 544.75,
-      "epoch": 0.0017708069820389577,
-      "grad_norm": 2.6467864513397217,
-      "kl": 0.000438690185546875,
-      "learning_rate": 9.991145965089805e-07,
-      "loss": 0.0,
-      "reward": 5.358240604400635,
-      "reward_std": 0.12582074105739594,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.7125557661056519,
-      "rewards/location_misalighment_heatmap_reward": 0.7356848120689392,
-      "rewards/overall_score_rewards": 0.9100000858306885,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 7
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 485.125,
-      "epoch": 0.002023779408044523,
-      "grad_norm": 2.7146477699279785,
-      "kl": 0.0003795623779296875,
-      "learning_rate": 9.989881102959778e-07,
-      "loss": 0.0,
-      "reward": 5.3003458976745605,
-      "reward_std": 0.12834414839744568,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.745768129825592,
-      "rewards/location_misalighment_heatmap_reward": 0.7758277058601379,
-      "rewards/overall_score_rewards": 0.7787500619888306,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 8
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 449.375,
-      "epoch": 0.0022767518340500886,
-      "grad_norm": 5.603612899780273,
-      "kl": 0.00057220458984375,
-      "learning_rate": 9.98861624082975e-07,
-      "loss": 0.0,
-      "reward": 5.459621429443359,
-      "reward_std": 0.17893511056900024,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.8214287757873535,
-      "rewards/location_misalighment_heatmap_reward": 0.8294430375099182,
-      "rewards/overall_score_rewards": 0.8087500333786011,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 9
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 483.25,
-      "epoch": 0.0025297242600556538,
-      "grad_norm": 2.887422561645508,
-      "kl": 0.0010986328125,
-      "learning_rate": 9.98735137869972e-07,
-      "loss": 0.0,
-      "reward": 5.473537445068359,
-      "reward_std": 0.0986822098493576,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.8078077435493469,
-      "rewards/location_misalighment_heatmap_reward": 0.8007298111915588,
-      "rewards/overall_score_rewards": 0.8649998903274536,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 10
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 436.5,
-      "epoch": 0.0027826966860612194,
-      "grad_norm": 3.3854877948760986,
-      "kl": 0.000461578369140625,
-      "learning_rate": 9.986086516569693e-07,
-      "loss": 0.0,
-      "reward": 5.272259712219238,
-      "reward_std": 0.11706135421991348,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.7914199829101562,
-      "rewards/location_misalighment_heatmap_reward": 0.6045900583267212,
-      "rewards/overall_score_rewards": 0.8762500286102295,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 11
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 514.25,
-      "epoch": 0.0030356691120667846,
-      "grad_norm": 3.44535231590271,
-      "kl": 0.00069427490234375,
-      "learning_rate": 9.984821654439667e-07,
-      "loss": 0.0,
-      "reward": 5.277377128601074,
-      "reward_std": 0.19741389155387878,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.7073001861572266,
-      "rewards/location_misalighment_heatmap_reward": 0.7613269090652466,
-      "rewards/overall_score_rewards": 0.8087499141693115,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 12
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 492.625,
-      "epoch": 0.0032886415380723503,
-      "grad_norm": 2.725538969039917,
-      "kl": 0.0004482269287109375,
-      "learning_rate": 9.983556792309638e-07,
-      "loss": 0.0,
-      "reward": 5.384651184082031,
-      "reward_std": 0.13790027797222137,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.7794829607009888,
-      "rewards/location_misalighment_heatmap_reward": 0.822668194770813,
-      "rewards/overall_score_rewards": 0.7825000286102295,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 13
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 501.125,
-      "epoch": 0.0035416139640779155,
-      "grad_norm": 5.83013916015625,
-      "kl": 0.00160980224609375,
-      "learning_rate": 9.98229193017961e-07,
-      "loss": 0.0,
-      "reward": 5.08870792388916,
-      "reward_std": 0.2164464145898819,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.5673180222511292,
-      "rewards/location_misalighment_heatmap_reward": 0.750139594078064,
-      "rewards/overall_score_rewards": 0.7712498903274536,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 14
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 484.25,
-      "epoch": 0.003794586390083481,
-      "grad_norm": 2.5159952640533447,
-      "kl": 0.000804901123046875,
-      "learning_rate": 9.981027068049582e-07,
-      "loss": 0.0,
-      "reward": 5.056214332580566,
-      "reward_std": 0.14230772852897644,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.6618410348892212,
-      "rewards/location_misalighment_heatmap_reward": 0.5518736243247986,
-      "rewards/overall_score_rewards": 0.8424999713897705,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 15
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 469.375,
-      "epoch": 0.004047558816089046,
-      "grad_norm": 6.992142677307129,
-      "kl": 0.00141143798828125,
-      "learning_rate": 9.979762205919555e-07,
-      "loss": 0.0001,
-      "reward": 5.003024578094482,
-      "reward_std": 0.8147953152656555,
-      "rewards/aesthetic_score_rewards": 0.875,
-      "rewards/location_artifact_heatmap_reward": 0.791554868221283,
-      "rewards/location_misalighment_heatmap_reward": 0.8352197408676147,
-      "rewards/overall_score_rewards": 0.7512500286102295,
-      "rewards/plausibility_score_rewards": 0.875,
-      "rewards/semantic_score_rewards": 0.875,
-      "step": 16
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 477.375,
-      "epoch": 0.0043005312420946115,
-      "grad_norm": 9.937246322631836,
-      "kl": 0.00115966796875,
-      "learning_rate": 9.978497343789528e-07,
-      "loss": 0.0,
-      "reward": 5.40995454788208,
-      "reward_std": 0.13958589732646942,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.7911748290061951,
-      "rewards/location_misalighment_heatmap_reward": 0.7762799263000488,
-      "rewards/overall_score_rewards": 0.8424999713897705,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 17
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 480.875,
-      "epoch": 0.004553503668100177,
-      "grad_norm": 3.1487274169921875,
-      "kl": 0.0012054443359375,
-      "learning_rate": 9.977232481659499e-07,
-      "loss": 0.0,
-      "reward": 5.3303937911987305,
-      "reward_std": 0.1731274574995041,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.7854811549186707,
-      "rewards/location_misalighment_heatmap_reward": 0.7286632061004639,
-      "rewards/overall_score_rewards": 0.8162500858306885,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 18
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 411.0,
-      "epoch": 0.004806476094105743,
-      "grad_norm": 5.541186332702637,
-      "kl": 0.000568389892578125,
-      "learning_rate": 9.97596761952947e-07,
-      "loss": 0.0,
-      "reward": 5.097373008728027,
-      "reward_std": 0.8139442205429077,
-      "rewards/aesthetic_score_rewards": 0.875,
-      "rewards/location_artifact_heatmap_reward": 0.9115674495697021,
-      "rewards/location_misalighment_heatmap_reward": 0.8095557689666748,
-      "rewards/overall_score_rewards": 0.7512500286102295,
-      "rewards/plausibility_score_rewards": 0.875,
-      "rewards/semantic_score_rewards": 0.875,
-      "step": 19
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 459.0,
-      "epoch": 0.0050594485201113076,
-      "grad_norm": 2.351588726043701,
-      "kl": 0.000637054443359375,
-      "learning_rate": 9.974702757399443e-07,
-      "loss": 0.0006,
-      "reward": 5.558006763458252,
-      "reward_std": 0.08841203153133392,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.722720205783844,
-      "rewards/location_misalighment_heatmap_reward": 0.8952864408493042,
-      "rewards/overall_score_rewards": 0.940000057220459,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 20
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 484.125,
-      "epoch": 0.005312420946116873,
-      "grad_norm": 3.0540034770965576,
-      "kl": 0.000675201416015625,
-      "learning_rate": 9.973437895269416e-07,
-      "loss": 0.0,
-      "reward": 5.327281951904297,
-      "reward_std": 0.08891787379980087,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.7103680968284607,
-      "rewards/location_misalighment_heatmap_reward": 0.8719138503074646,
-      "rewards/overall_score_rewards": 0.7450000047683716,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 21
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 466.0,
-      "epoch": 0.005565393372122439,
-      "grad_norm": 2.333858013153076,
-      "kl": 0.00052642822265625,
-      "learning_rate": 9.972173033139387e-07,
-      "loss": 0.0006,
-      "reward": 5.494153022766113,
-      "reward_std": 0.04288619011640549,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.8250214457511902,
-      "rewards/location_misalighment_heatmap_reward": 0.8191317915916443,
-      "rewards/overall_score_rewards": 0.8500000834465027,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 22
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 486.875,
-      "epoch": 0.0058183657981280045,
-      "grad_norm": 3.2977066040039062,
-      "kl": 0.001800537109375,
-      "learning_rate": 9.970908171009358e-07,
-      "loss": 0.0001,
-      "reward": 5.396883487701416,
-      "reward_std": 0.14189015328884125,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.7350619435310364,
-      "rewards/location_misalighment_heatmap_reward": 0.7555716037750244,
-      "rewards/overall_score_rewards": 0.90625,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 23
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 498.625,
-      "epoch": 0.006071338224133569,
-      "grad_norm": 3.383419990539551,
-      "kl": 0.0020904541015625,
-      "learning_rate": 9.969643308879331e-07,
-      "loss": 0.0001,
-      "reward": 5.401487827301025,
-      "reward_std": 0.1331316977739334,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.6623430252075195,
-      "rewards/location_misalighment_heatmap_reward": 0.8328946828842163,
-      "rewards/overall_score_rewards": 0.90625,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 24
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 452.0,
-      "epoch": 0.006324310650139135,
-      "grad_norm": 2.778303861618042,
-      "kl": 0.00112152099609375,
-      "learning_rate": 9.968378446749304e-07,
-      "loss": 0.0,
-      "reward": 5.541274070739746,
-      "reward_std": 0.13871176540851593,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.8679863214492798,
-      "rewards/location_misalighment_heatmap_reward": 0.7632877826690674,
-      "rewards/overall_score_rewards": 0.9100000262260437,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 25
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 464.5,
-      "epoch": 0.0065772830761447005,
-      "grad_norm": 5.355808258056641,
-      "kl": 0.00128173828125,
-      "learning_rate": 9.967113584619276e-07,
-      "loss": 0.0001,
-      "reward": 5.467164993286133,
-      "reward_std": 0.10859573632478714,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.6112275719642639,
-      "rewards/location_misalighment_heatmap_reward": 0.9196872711181641,
-      "rewards/overall_score_rewards": 0.9362499713897705,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 26
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 458.0,
-      "epoch": 0.006830255502150265,
-      "grad_norm": 2.5720341205596924,
-      "kl": 0.00131988525390625,
-      "learning_rate": 9.965848722489249e-07,
-      "loss": 0.0001,
-      "reward": 5.522697925567627,
-      "reward_std": 0.042607516050338745,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.7024801969528198,
-      "rewards/location_misalighment_heatmap_reward": 0.9102176427841187,
-      "rewards/overall_score_rewards": 0.9100000858306885,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 27
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 480.875,
-      "epoch": 0.007083227928155831,
-      "grad_norm": 3.569387435913086,
-      "kl": 0.00165557861328125,
-      "learning_rate": 9.96458386035922e-07,
-      "loss": 0.0003,
-      "reward": 5.367475986480713,
-      "reward_std": 0.10307542979717255,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.7223396301269531,
-      "rewards/location_misalighment_heatmap_reward": 0.7726362943649292,
-      "rewards/overall_score_rewards": 0.872499942779541,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 28
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 470.625,
-      "epoch": 0.007336200354161397,
-      "grad_norm": 5.62157678604126,
-      "kl": 0.0023193359375,
-      "learning_rate": 9.963318998229193e-07,
-      "loss": 0.0001,
-      "reward": 5.359262943267822,
-      "reward_std": 0.10050362348556519,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.8168458938598633,
-      "rewards/location_misalighment_heatmap_reward": 0.6661669611930847,
-      "rewards/overall_score_rewards": 0.8762500286102295,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 29
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 494.875,
-      "epoch": 0.007589172780166962,
-      "grad_norm": 2.860589027404785,
-      "kl": 0.00194549560546875,
-      "learning_rate": 9.962054136099166e-07,
-      "loss": 0.0001,
-      "reward": 5.405731678009033,
-      "reward_std": 0.1737304925918579,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.7799570560455322,
-      "rewards/location_misalighment_heatmap_reward": 0.8132746815681458,
-      "rewards/overall_score_rewards": 0.8125,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 30
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 437.125,
-      "epoch": 0.007842145206172527,
-      "grad_norm": 3.621122360229492,
-      "kl": 0.00439453125,
-      "learning_rate": 9.960789273969137e-07,
-      "loss": 0.0002,
-      "reward": 4.794467926025391,
-      "reward_std": 0.7571890950202942,
-      "rewards/aesthetic_score_rewards": 0.875,
-      "rewards/location_artifact_heatmap_reward": 0.7693961262702942,
-      "rewards/location_misalighment_heatmap_reward": 0.6788214445114136,
-      "rewards/overall_score_rewards": 0.721250057220459,
-      "rewards/plausibility_score_rewards": 0.875,
-      "rewards/semantic_score_rewards": 0.875,
-      "step": 31
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 485.0,
-      "epoch": 0.008095117632178092,
-      "grad_norm": 5.9346489906311035,
-      "kl": 0.00482177734375,
-      "learning_rate": 9.959524411839108e-07,
-      "loss": 0.0002,
-      "reward": 5.404839515686035,
-      "reward_std": 0.1227060854434967,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.7918017506599426,
-      "rewards/location_misalighment_heatmap_reward": 0.8305379152297974,
-      "rewards/overall_score_rewards": 0.7825000286102295,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 32
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 410.25,
-      "epoch": 0.008348090058183658,
-      "grad_norm": 4.121129512786865,
-      "kl": 0.017822265625,
-      "learning_rate": 9.958259549709081e-07,
-      "loss": 0.0007,
-      "reward": 5.0262908935546875,
-      "reward_std": 0.7571164965629578,
-      "rewards/aesthetic_score_rewards": 0.875,
-      "rewards/location_artifact_heatmap_reward": 0.8097456693649292,
-      "rewards/location_misalighment_heatmap_reward": 0.8702950477600098,
-      "rewards/overall_score_rewards": 0.721250057220459,
-      "rewards/plausibility_score_rewards": 0.875,
-      "rewards/semantic_score_rewards": 0.875,
-      "step": 33
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 494.125,
-      "epoch": 0.008601062484189223,
-      "grad_norm": 68.54109954833984,
-      "kl": 0.006256103515625,
-      "learning_rate": 9.956994687579054e-07,
-      "loss": 0.0002,
-      "reward": 5.404845237731934,
-      "reward_std": 0.09779535233974457,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.6971195936203003,
-      "rewards/location_misalighment_heatmap_reward": 0.8914761543273926,
-      "rewards/overall_score_rewards": 0.8162500858306885,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 34
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 479.625,
-      "epoch": 0.00885403491019479,
-      "grad_norm": 6.452138423919678,
-      "kl": 0.011474609375,
-      "learning_rate": 9.955729825449025e-07,
-      "loss": 0.0005,
-      "reward": 5.1309075355529785,
-      "reward_std": 0.16118118166923523,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.5963692665100098,
-      "rewards/location_misalighment_heatmap_reward": 0.7182881832122803,
-      "rewards/overall_score_rewards": 0.8162500858306885,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 35
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 425.25,
-      "epoch": 0.009107007336200354,
-      "grad_norm": 3.74330997467041,
-      "kl": 0.01708984375,
-      "learning_rate": 9.954464963318998e-07,
-      "loss": 0.0007,
-      "reward": 4.736416339874268,
-      "reward_std": 0.8011575937271118,
-      "rewards/aesthetic_score_rewards": 0.875,
-      "rewards/location_artifact_heatmap_reward": 0.6437466144561768,
-      "rewards/location_misalighment_heatmap_reward": 0.7501698732376099,
-      "rewards/overall_score_rewards": 0.7175000309944153,
-      "rewards/plausibility_score_rewards": 0.875,
-      "rewards/semantic_score_rewards": 0.875,
-      "step": 36
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 469.875,
-      "epoch": 0.009359979762205919,
-      "grad_norm": 5.57287073135376,
-      "kl": 0.0234375,
-      "learning_rate": 9.95320010118897e-07,
-      "loss": 0.0009,
-      "reward": 5.053951263427734,
-      "reward_std": 0.09248366951942444,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.5215615034103394,
-      "rewards/location_misalighment_heatmap_reward": 0.6523897051811218,
-      "rewards/overall_score_rewards": 0.8799999952316284,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 37
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 474.5,
-      "epoch": 0.009612952188211486,
-      "grad_norm": 2.297215223312378,
-      "kl": 0.01214599609375,
-      "learning_rate": 9.951935239058942e-07,
-      "loss": 0.0005,
-      "reward": 4.92038631439209,
-      "reward_std": 0.31508731842041016,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.6160138249397278,
-      "rewards/location_misalighment_heatmap_reward": 0.7093722820281982,
-      "rewards/overall_score_rewards": 0.5950000286102295,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 38
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 426.25,
-      "epoch": 0.00986592461421705,
-      "grad_norm": 7.791620254516602,
-      "kl": 0.02880859375,
-      "learning_rate": 9.950670376928916e-07,
-      "loss": 0.0011,
-      "reward": 4.909842491149902,
-      "reward_std": 0.7547658681869507,
-      "rewards/aesthetic_score_rewards": 0.875,
-      "rewards/location_artifact_heatmap_reward": 0.7402044534683228,
-      "rewards/location_misalighment_heatmap_reward": 0.8271378874778748,
-      "rewards/overall_score_rewards": 0.7174999713897705,
-      "rewards/plausibility_score_rewards": 0.875,
-      "rewards/semantic_score_rewards": 0.875,
-      "step": 39
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 446.875,
-      "epoch": 0.010118897040222615,
-      "grad_norm": 5.844668865203857,
-      "kl": 0.0225830078125,
-      "learning_rate": 9.949405514798887e-07,
-      "loss": 0.0009,
-      "reward": 5.615790367126465,
-      "reward_std": 0.05132600665092468,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.914071261882782,
-      "rewards/location_misalighment_heatmap_reward": 0.8517189025878906,
-      "rewards/overall_score_rewards": 0.8500001430511475,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 40
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 482.75,
-      "epoch": 0.010371869466228182,
-      "grad_norm": 2.13356351852417,
-      "kl": 0.01318359375,
-      "learning_rate": 9.948140652668858e-07,
-      "loss": 0.0008,
-      "reward": 5.337843894958496,
-      "reward_std": 0.05293720215559006,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.7141811847686768,
-      "rewards/location_misalighment_heatmap_reward": 0.6574122905731201,
-      "rewards/overall_score_rewards": 0.966249942779541,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 41
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 468.625,
-      "epoch": 0.010624841892233746,
-      "grad_norm": 3.4011526107788086,
-      "kl": 0.0260009765625,
-      "learning_rate": 9.94687579053883e-07,
-      "loss": 0.001,
-      "reward": 5.161016464233398,
-      "reward_std": 0.09107545763254166,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.4952339231967926,
-      "rewards/location_misalighment_heatmap_reward": 0.7295328974723816,
-      "rewards/overall_score_rewards": 0.9362499713897705,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 42
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 375.25,
-      "epoch": 0.010877814318239311,
-      "grad_norm": 5.546986103057861,
-      "kl": 0.035888671875,
-      "learning_rate": 9.945610928408804e-07,
-      "loss": 0.0014,
-      "reward": 5.538779258728027,
-      "reward_std": 0.11370107531547546,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.824926495552063,
-      "rewards/location_misalighment_heatmap_reward": 0.8038531541824341,
-      "rewards/overall_score_rewards": 0.9100000858306885,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 43
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 431.625,
-      "epoch": 0.011130786744244878,
-      "grad_norm": 20.80287742614746,
-      "kl": 0.03173828125,
-      "learning_rate": 9.944346066278775e-07,
-      "loss": 0.0014,
-      "reward": 5.45789909362793,
-      "reward_std": 0.041689638048410416,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.6997177600860596,
-      "rewards/location_misalighment_heatmap_reward": 0.8481816053390503,
-      "rewards/overall_score_rewards": 0.9100000262260437,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 44
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 434.25,
-      "epoch": 0.011383759170250442,
-      "grad_norm": 8.854325294494629,
-      "kl": 0.02490234375,
-      "learning_rate": 9.943081204148748e-07,
-      "loss": 0.001,
-      "reward": 5.107644557952881,
-      "reward_std": 0.15394118428230286,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.6861124038696289,
-      "rewards/location_misalighment_heatmap_reward": 0.7702821493148804,
-      "rewards/overall_score_rewards": 0.6512500047683716,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 45
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 372.75,
-      "epoch": 0.011636731596256009,
-      "grad_norm": 6.125293731689453,
-      "kl": 0.043701171875,
-      "learning_rate": 9.94181634201872e-07,
-      "loss": 0.0017,
-      "reward": 4.9134674072265625,
-      "reward_std": 0.7654469609260559,
-      "rewards/aesthetic_score_rewards": 0.875,
-      "rewards/location_artifact_heatmap_reward": 0.7635915279388428,
-      "rewards/location_misalighment_heatmap_reward": 0.8711254596710205,
-      "rewards/overall_score_rewards": 0.6537500023841858,
-      "rewards/plausibility_score_rewards": 0.875,
-      "rewards/semantic_score_rewards": 0.875,
-      "step": 46
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 456.375,
-      "epoch": 0.011889704022261574,
-      "grad_norm": 3.2194080352783203,
-      "kl": 0.0206298828125,
-      "learning_rate": 9.940551479888692e-07,
-      "loss": 0.0008,
-      "reward": 5.286049842834473,
-      "reward_std": 0.09385614097118378,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.7511510848999023,
-      "rewards/location_misalighment_heatmap_reward": 0.7861489057540894,
-      "rewards/overall_score_rewards": 0.7487500309944153,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 47
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 409.375,
-      "epoch": 0.012142676448267139,
-      "grad_norm": 4.2013325691223145,
-      "kl": 0.03759765625,
-      "learning_rate": 9.939286617758665e-07,
-      "loss": 0.0015,
-      "reward": 5.483973503112793,
-      "reward_std": 0.09291357547044754,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.8019596934318542,
-      "rewards/location_misalighment_heatmap_reward": 0.8207635879516602,
-      "rewards/overall_score_rewards": 0.8612500429153442,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 48
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 419.375,
-      "epoch": 0.012395648874272705,
-      "grad_norm": 4.354659080505371,
-      "kl": 0.02734375,
-      "learning_rate": 9.938021755628636e-07,
-      "loss": 0.0011,
-      "reward": 4.686689853668213,
-      "reward_std": 0.4714539647102356,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.6772425174713135,
-      "rewards/location_misalighment_heatmap_reward": 0.643197238445282,
-      "rewards/overall_score_rewards": 0.36625003814697266,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 49
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 404.625,
-      "epoch": 0.01264862130027827,
-      "grad_norm": 3.6461737155914307,
-      "kl": 0.03173828125,
-      "learning_rate": 9.936756893498607e-07,
-      "loss": 0.0013,
-      "reward": 5.545103073120117,
-      "reward_std": 0.06625383347272873,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.795859158039093,
-      "rewards/location_misalighment_heatmap_reward": 0.7829937934875488,
-      "rewards/overall_score_rewards": 0.966249942779541,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 50
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 451.125,
-      "epoch": 0.012901593726283835,
-      "grad_norm": 3.7185897827148438,
-      "kl": 0.0206298828125,
-      "learning_rate": 9.93549203136858e-07,
-      "loss": 0.0008,
-      "reward": 5.4477219581604,
-      "reward_std": 0.0937633365392685,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.7628577947616577,
-      "rewards/location_misalighment_heatmap_reward": 0.808614194393158,
-      "rewards/overall_score_rewards": 0.8762500286102295,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 51
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 433.375,
-      "epoch": 0.013154566152289401,
-      "grad_norm": 3.265334129333496,
-      "kl": 0.0186767578125,
-      "learning_rate": 9.934227169238554e-07,
-      "loss": 0.0007,
-      "reward": 5.501402854919434,
-      "reward_std": 0.11307864636182785,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.6259495615959167,
-      "rewards/location_misalighment_heatmap_reward": 0.9392036199569702,
-      "rewards/overall_score_rewards": 0.9362499713897705,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 52
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 513.25,
-      "epoch": 0.013407538578294966,
-      "grad_norm": 3.123044013977051,
-      "kl": 0.01300048828125,
-      "learning_rate": 9.932962307108525e-07,
-      "loss": 0.0005,
-      "reward": 5.397100448608398,
-      "reward_std": 0.14120273292064667,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.8524265885353088,
-      "rewards/location_misalighment_heatmap_reward": 0.7659234404563904,
-      "rewards/overall_score_rewards": 0.778749942779541,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 53
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 407.875,
-      "epoch": 0.01366051100430053,
-      "grad_norm": 6.148492336273193,
-      "kl": 0.023681640625,
-      "learning_rate": 9.931697444978498e-07,
-      "loss": 0.001,
-      "reward": 5.381246566772461,
-      "reward_std": 0.17179468274116516,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.7448151111602783,
-      "rewards/location_misalighment_heatmap_reward": 0.8201813697814941,
-      "rewards/overall_score_rewards": 0.8162500262260437,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 54
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 303.125,
-      "epoch": 0.013913483430306097,
-      "grad_norm": 5.236420154571533,
-      "kl": 0.060302734375,
-      "learning_rate": 9.930432582848469e-07,
-      "loss": 0.003,
-      "reward": 5.038935661315918,
-      "reward_std": 0.2707816958427429,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.6911619901657104,
-      "rewards/location_misalighment_heatmap_reward": 0.7827734351158142,
-      "rewards/overall_score_rewards": 0.5649999976158142,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 55
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 418.5,
-      "epoch": 0.014166455856311662,
-      "grad_norm": 6.169740200042725,
-      "kl": 0.037841796875,
-      "learning_rate": 9.929167720718442e-07,
-      "loss": 0.0015,
-      "reward": 5.141636848449707,
-      "reward_std": 0.00907151959836483,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.6578080654144287,
-      "rewards/location_misalighment_heatmap_reward": 0.6713287234306335,
-      "rewards/overall_score_rewards": 0.8125000596046448,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 56
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 430.75,
-      "epoch": 0.014419428282317227,
-      "grad_norm": 3.0601601600646973,
-      "kl": 0.0186767578125,
-      "learning_rate": 9.927902858588413e-07,
-      "loss": 0.0008,
-      "reward": 5.162736892700195,
-      "reward_std": 0.01670549251139164,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.6311650276184082,
-      "rewards/location_misalighment_heatmap_reward": 0.6515717506408691,
-      "rewards/overall_score_rewards": 0.8800000548362732,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 57
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 320.25,
-      "epoch": 0.014672400708322793,
-      "grad_norm": 7.362841606140137,
-      "kl": 0.0615234375,
-      "learning_rate": 9.926637996458386e-07,
-      "loss": 0.0025,
-      "reward": 4.972415447235107,
-      "reward_std": 0.2225758284330368,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.6381765604019165,
-      "rewards/location_misalighment_heatmap_reward": 0.8667387366294861,
-      "rewards/overall_score_rewards": 0.4675000309944153,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 58
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 321.375,
-      "epoch": 0.014925373134328358,
-      "grad_norm": 9.785571098327637,
-      "kl": 0.06494140625,
-      "learning_rate": 9.925373134328357e-07,
-      "loss": 0.0026,
-      "reward": 5.335415840148926,
-      "reward_std": 0.08124318718910217,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.6716148853302002,
-      "rewards/location_misalighment_heatmap_reward": 0.8438011407852173,
-      "rewards/overall_score_rewards": 0.8200001120567322,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 59
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 435.25,
-      "epoch": 0.015178345560333924,
-      "grad_norm": 2.9662091732025146,
-      "kl": 0.004852294921875,
-      "learning_rate": 9.92410827219833e-07,
-      "loss": 0.0002,
-      "reward": 5.437358856201172,
-      "reward_std": 0.049164045602083206,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.6008865833282471,
-      "rewards/location_misalighment_heatmap_reward": 0.8664723634719849,
-      "rewards/overall_score_rewards": 0.9700000286102295,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 60
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 471.5,
-      "epoch": 0.01543131798633949,
-      "grad_norm": 2.7331151962280273,
-      "kl": 0.0037994384765625,
-      "learning_rate": 9.922843410068303e-07,
-      "loss": 0.0001,
-      "reward": 5.221156120300293,
-      "reward_std": 0.05214188992977142,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.46495527029037476,
-      "rewards/location_misalighment_heatmap_reward": 0.8499510288238525,
-      "rewards/overall_score_rewards": 0.90625,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 61
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 511.875,
-      "epoch": 0.015684290412345054,
-      "grad_norm": 3.0238466262817383,
-      "kl": 0.00127410888671875,
-      "learning_rate": 9.921578547938274e-07,
-      "loss": 0.0001,
-      "reward": 5.261039733886719,
-      "reward_std": 0.13529786467552185,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.6692686676979065,
-      "rewards/location_misalighment_heatmap_reward": 0.8092712163925171,
-      "rewards/overall_score_rewards": 0.7825000286102295,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 62
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 476.625,
-      "epoch": 0.01593726283835062,
-      "grad_norm": 5.247809410095215,
-      "kl": 0.001953125,
-      "learning_rate": 9.920313685808245e-07,
-      "loss": 0.0001,
-      "reward": 5.164344787597656,
-      "reward_std": 0.10323432087898254,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.34531018137931824,
-      "rewards/location_misalighment_heatmap_reward": 0.9427845478057861,
-      "rewards/overall_score_rewards": 0.8762500286102295,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 63
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 528.75,
-      "epoch": 0.016190235264356184,
-      "grad_norm": 15.859211921691895,
-      "kl": 0.0027313232421875,
-      "learning_rate": 9.919048823678218e-07,
-      "loss": 0.0001,
-      "reward": 5.229736328125,
-      "reward_std": 0.08716608583927155,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.7767056226730347,
-      "rewards/location_misalighment_heatmap_reward": 0.7005307078361511,
-      "rewards/overall_score_rewards": 0.752500057220459,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 64
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 466.5,
-      "epoch": 0.016443207690361752,
-      "grad_norm": 3.4766626358032227,
-      "kl": 0.0013580322265625,
-      "learning_rate": 9.917783961548192e-07,
-      "loss": 0.0001,
-      "reward": 5.607626914978027,
-      "reward_std": 0.12780779600143433,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.7906875610351562,
-      "rewards/location_misalighment_heatmap_reward": 0.9669394493103027,
-      "rewards/overall_score_rewards": 0.8500000834465027,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 65
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 458.375,
-      "epoch": 0.016696180116367317,
-      "grad_norm": 4.968540191650391,
-      "kl": 0.00372314453125,
-      "learning_rate": 9.916519099418163e-07,
-      "loss": 0.0001,
-      "reward": 5.4170660972595215,
-      "reward_std": 0.1274934560060501,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.7114583253860474,
-      "rewards/location_misalighment_heatmap_reward": 0.7993578910827637,
-      "rewards/overall_score_rewards": 0.90625,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 66
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 484.875,
-      "epoch": 0.01694915254237288,
-      "grad_norm": 4.454622268676758,
-      "kl": 0.0025634765625,
-      "learning_rate": 9.915254237288136e-07,
-      "loss": 0.0001,
-      "reward": 5.413036346435547,
-      "reward_std": 0.055481187999248505,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.8084956407546997,
-      "rewards/location_misalighment_heatmap_reward": 0.7620410919189453,
-      "rewards/overall_score_rewards": 0.8424999713897705,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 67
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 521.125,
-      "epoch": 0.017202124968378446,
-      "grad_norm": 2.575721263885498,
-      "kl": 0.001861572265625,
-      "learning_rate": 9.913989375158107e-07,
-      "loss": 0.0001,
-      "reward": 5.65494441986084,
-      "reward_std": 0.09674935042858124,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.885114312171936,
-      "rewards/location_misalighment_heatmap_reward": 0.8710803985595703,
-      "rewards/overall_score_rewards": 0.8987499475479126,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 68
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 472.875,
-      "epoch": 0.01745509739438401,
-      "grad_norm": 2.5332581996917725,
-      "kl": 0.00130462646484375,
-      "learning_rate": 9.91272451302808e-07,
-      "loss": 0.0001,
-      "reward": 5.421402454376221,
-      "reward_std": 0.18434956669807434,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.9000935554504395,
-      "rewards/location_misalighment_heatmap_reward": 0.7688089609146118,
-      "rewards/overall_score_rewards": 0.752500057220459,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 69
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 466.5,
-      "epoch": 0.01770806982038958,
-      "grad_norm": 5.255011558532715,
-      "kl": 0.0054931640625,
-      "learning_rate": 9.911459650898053e-07,
-      "loss": 0.0002,
-      "reward": 4.905242443084717,
-      "reward_std": 0.17114639282226562,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.4978741407394409,
-      "rewards/location_misalighment_heatmap_reward": 0.5911184549331665,
-      "rewards/overall_score_rewards": 0.8162500858306885,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 70
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 488.75,
-      "epoch": 0.017961042246395144,
-      "grad_norm": 2.3159234523773193,
-      "kl": 0.00160980224609375,
-      "learning_rate": 9.910194788768024e-07,
-      "loss": 0.0,
-      "reward": 5.508110046386719,
-      "reward_std": 0.04277314990758896,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.8036966323852539,
-      "rewards/location_misalighment_heatmap_reward": 0.8544133305549622,
-      "rewards/overall_score_rewards": 0.8500000834465027,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 71
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 483.125,
-      "epoch": 0.01821401467240071,
-      "grad_norm": 4.344380855560303,
-      "kl": 0.0016632080078125,
-      "learning_rate": 9.908929926637995e-07,
-      "loss": 0.0001,
-      "reward": 5.345328330993652,
-      "reward_std": 0.09193618595600128,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.7422531843185425,
-      "rewards/location_misalighment_heatmap_reward": 0.7868252396583557,
-      "rewards/overall_score_rewards": 0.8162499666213989,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 72
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 472.625,
-      "epoch": 0.018466987098406273,
-      "grad_norm": 3.2257914543151855,
-      "kl": 0.00144195556640625,
-      "learning_rate": 9.907665064507968e-07,
-      "loss": 0.0005,
-      "reward": 5.4638800621032715,
-      "reward_std": 0.09152112156152725,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.6981525421142578,
-      "rewards/location_misalighment_heatmap_reward": 0.8894773721694946,
-      "rewards/overall_score_rewards": 0.8762500286102295,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 73
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 503.25,
-      "epoch": 0.018719959524411838,
-      "grad_norm": 2.601001501083374,
-      "kl": 0.003662109375,
-      "learning_rate": 9.906400202377941e-07,
-      "loss": 0.0001,
-      "reward": 5.326350688934326,
-      "reward_std": 0.08632879704236984,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.6927300095558167,
-      "rewards/location_misalighment_heatmap_reward": 0.7536205649375916,
-      "rewards/overall_score_rewards": 0.8800000548362732,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 74
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 486.125,
-      "epoch": 0.018972931950417403,
-      "grad_norm": 4.618655204772949,
-      "kl": 0.00168609619140625,
-      "learning_rate": 9.905135340247912e-07,
-      "loss": 0.0001,
-      "reward": 5.580634117126465,
-      "reward_std": 0.14603908360004425,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.8971958160400391,
-      "rewards/location_misalighment_heatmap_reward": 0.8446887135505676,
-      "rewards/overall_score_rewards": 0.8387500047683716,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 75
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 507.375,
-      "epoch": 0.01922590437642297,
-      "grad_norm": 3.295837640762329,
-      "kl": 0.00177001953125,
-      "learning_rate": 9.903870478117885e-07,
-      "loss": 0.0001,
-      "reward": 5.342641830444336,
-      "reward_std": 0.10630211234092712,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.7530924081802368,
-      "rewards/location_misalighment_heatmap_reward": 0.6832996010780334,
-      "rewards/overall_score_rewards": 0.90625,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 76
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 477.25,
-      "epoch": 0.019478876802428536,
-      "grad_norm": 5.158864498138428,
-      "kl": 0.00182342529296875,
-      "learning_rate": 9.902605615987856e-07,
-      "loss": 0.0001,
-      "reward": 5.472329616546631,
-      "reward_std": 0.20435674488544464,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.7965782880783081,
-      "rewards/location_misalighment_heatmap_reward": 0.8632513284683228,
-      "rewards/overall_score_rewards": 0.8125,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 77
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 458.0,
-      "epoch": 0.0197318492284341,
-      "grad_norm": 3.044123649597168,
-      "kl": 0.00634765625,
-      "learning_rate": 9.90134075385783e-07,
-      "loss": 0.0003,
-      "reward": 5.425571918487549,
-      "reward_std": 0.09147030115127563,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.8515055179595947,
-      "rewards/location_misalighment_heatmap_reward": 0.6940664052963257,
-      "rewards/overall_score_rewards": 0.8800000548362732,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 78
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 471.625,
-      "epoch": 0.019984821654439665,
-      "grad_norm": 5.804065227508545,
-      "kl": 0.00201416015625,
-      "learning_rate": 9.9000758917278e-07,
-      "loss": 0.0001,
-      "reward": 5.500574111938477,
-      "reward_std": 0.09330915659666061,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.7922377586364746,
-      "rewards/location_misalighment_heatmap_reward": 0.8920866847038269,
-      "rewards/overall_score_rewards": 0.8162499666213989,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 79
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 500.125,
-      "epoch": 0.02023779408044523,
-      "grad_norm": 7.658450126647949,
-      "kl": 0.001953125,
-      "learning_rate": 9.898811029597774e-07,
-      "loss": 0.0001,
-      "reward": 5.096043586730957,
-      "reward_std": 0.10247534513473511,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.5063791275024414,
-      "rewards/location_misalighment_heatmap_reward": 0.8409141302108765,
-      "rewards/overall_score_rewards": 0.7487500309944153,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 80
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 452.25,
-      "epoch": 0.0204907665064508,
-      "grad_norm": 5.492777347564697,
-      "kl": 0.00182342529296875,
-      "learning_rate": 9.897546167467745e-07,
-      "loss": 0.0001,
-      "reward": 5.383922576904297,
-      "reward_std": 0.09438686072826385,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.7550057172775269,
-      "rewards/location_misalighment_heatmap_reward": 0.7526668310165405,
-      "rewards/overall_score_rewards": 0.8762500286102295,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 81
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 450.125,
-      "epoch": 0.020743738932456363,
-      "grad_norm": 5.516570568084717,
-      "kl": 0.00160980224609375,
-      "learning_rate": 9.896281305337718e-07,
-      "loss": 0.0001,
-      "reward": 5.369076251983643,
-      "reward_std": 0.12602126598358154,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.8047648668289185,
-      "rewards/location_misalighment_heatmap_reward": 0.7593116164207458,
-      "rewards/overall_score_rewards": 0.8050000071525574,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 82
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 491.5,
-      "epoch": 0.020996711358461928,
-      "grad_norm": 3.835981607437134,
-      "kl": 0.001800537109375,
-      "learning_rate": 9.89501644320769e-07,
-      "loss": 0.0001,
-      "reward": 5.489201545715332,
-      "reward_std": 0.08540283143520355,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.7125409245491028,
-      "rewards/location_misalighment_heatmap_reward": 0.9004106521606445,
-      "rewards/overall_score_rewards": 0.8762500286102295,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 83
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 496.375,
-      "epoch": 0.021249683784467493,
-      "grad_norm": 2.8858718872070312,
-      "kl": 0.002593994140625,
-      "learning_rate": 9.893751581077662e-07,
-      "loss": 0.0001,
-      "reward": 5.359023571014404,
-      "reward_std": 0.1377764344215393,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.7235873937606812,
-      "rewards/location_misalighment_heatmap_reward": 0.8529358506202698,
-      "rewards/overall_score_rewards": 0.7825000286102295,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 84
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 475.125,
-      "epoch": 0.021502656210473058,
-      "grad_norm": 2.8391635417938232,
-      "kl": 0.00176239013671875,
-      "learning_rate": 9.892486718947635e-07,
-      "loss": 0.0001,
-      "reward": 5.476465225219727,
-      "reward_std": 0.05020717903971672,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.8874689936637878,
-      "rewards/location_misalighment_heatmap_reward": 0.8102465271949768,
-      "rewards/overall_score_rewards": 0.7787500023841858,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 85
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 483.625,
-      "epoch": 0.021755628636478622,
-      "grad_norm": 3.5714197158813477,
-      "kl": 0.00225830078125,
-      "learning_rate": 9.891221856817606e-07,
-      "loss": 0.0001,
-      "reward": 5.499442100524902,
-      "reward_std": 0.1941453218460083,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.8186551332473755,
-      "rewards/location_misalighment_heatmap_reward": 0.8682875037193298,
-      "rewards/overall_score_rewards": 0.8125,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 86
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 463.0,
-      "epoch": 0.02200860106248419,
-      "grad_norm": 3.722184658050537,
-      "kl": 0.00185394287109375,
-      "learning_rate": 9.88995699468758e-07,
-      "loss": 0.0001,
-      "reward": 5.329732894897461,
-      "reward_std": 0.14551502466201782,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.7463504076004028,
-      "rewards/location_misalighment_heatmap_reward": 0.7371325492858887,
-      "rewards/overall_score_rewards": 0.846250057220459,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 87
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 481.5,
-      "epoch": 0.022261573488489755,
-      "grad_norm": 6.236143589019775,
-      "kl": 0.00238037109375,
-      "learning_rate": 9.88869213255755e-07,
-      "loss": 0.0001,
-      "reward": 5.474348068237305,
-      "reward_std": 0.09561501443386078,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.729633629322052,
-      "rewards/location_misalighment_heatmap_reward": 0.8684642910957336,
-      "rewards/overall_score_rewards": 0.8762500286102295,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 88
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 455.375,
-      "epoch": 0.02251454591449532,
-      "grad_norm": 2.780856132507324,
-      "kl": 0.00173187255859375,
-      "learning_rate": 9.887427270427523e-07,
-      "loss": 0.0001,
-      "reward": 5.335257053375244,
-      "reward_std": 0.10799127072095871,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.7437273263931274,
-      "rewards/location_misalighment_heatmap_reward": 0.779029905796051,
-      "rewards/overall_score_rewards": 0.8125,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 89
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 448.125,
-      "epoch": 0.022767518340500885,
-      "grad_norm": 4.681956768035889,
-      "kl": 0.001678466796875,
-      "learning_rate": 9.886162408297494e-07,
-      "loss": 0.0001,
-      "reward": 5.363909721374512,
-      "reward_std": 0.22407092154026031,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.6742722392082214,
-      "rewards/location_misalighment_heatmap_reward": 0.9108874797821045,
-      "rewards/overall_score_rewards": 0.778749942779541,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 90
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 505.125,
-      "epoch": 0.02302049076650645,
-      "grad_norm": 10.409021377563477,
-      "kl": 0.00347900390625,
-      "learning_rate": 9.884897546167468e-07,
-      "loss": 0.0001,
-      "reward": 5.44618034362793,
-      "reward_std": 0.039001766592264175,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.8661985397338867,
-      "rewards/location_misalighment_heatmap_reward": 0.8612315654754639,
-      "rewards/overall_score_rewards": 0.71875,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 91
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 493.25,
-      "epoch": 0.023273463192512018,
-      "grad_norm": 3.5002048015594482,
-      "kl": 0.0023345947265625,
-      "learning_rate": 9.883632684037439e-07,
-      "loss": 0.0003,
-      "reward": 5.219611167907715,
-      "reward_std": 0.0013880643527954817,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.6183225512504578,
-      "rewards/location_misalighment_heatmap_reward": 0.7887887954711914,
-      "rewards/overall_score_rewards": 0.8125,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 92
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 469.125,
-      "epoch": 0.023526435618517583,
-      "grad_norm": 11.687639236450195,
-      "kl": 0.0022430419921875,
-      "learning_rate": 9.882367821907412e-07,
-      "loss": 0.0001,
-      "reward": 5.291659355163574,
-      "reward_std": 0.13471651077270508,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.6197956800460815,
-      "rewards/location_misalighment_heatmap_reward": 0.8256136178970337,
-      "rewards/overall_score_rewards": 0.846250057220459,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 93
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 483.125,
-      "epoch": 0.023779408044523147,
-      "grad_norm": 10.824652671813965,
-      "kl": 0.002227783203125,
-      "learning_rate": 9.881102959777385e-07,
-      "loss": 0.0001,
-      "reward": 5.5981974601745605,
-      "reward_std": 0.13304394483566284,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.8992058634757996,
-      "rewards/location_misalighment_heatmap_reward": 0.848991870880127,
-      "rewards/overall_score_rewards": 0.8500000834465027,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 94
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 515.625,
-      "epoch": 0.024032380470528712,
-      "grad_norm": 3.047670364379883,
-      "kl": 0.001983642578125,
-      "learning_rate": 9.879838097647356e-07,
-      "loss": 0.0001,
-      "reward": 5.459529876708984,
-      "reward_std": 0.09775835275650024,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.6944575309753418,
-      "rewards/location_misalighment_heatmap_reward": 0.8325725197792053,
-      "rewards/overall_score_rewards": 0.9324999451637268,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 95
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 461.625,
-      "epoch": 0.024285352896534277,
-      "grad_norm": 5.527977466583252,
-      "kl": 0.0020599365234375,
-      "learning_rate": 9.878573235517329e-07,
-      "loss": 0.0001,
-      "reward": 5.455284595489502,
-      "reward_std": 0.10044293105602264,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.8116781711578369,
-      "rewards/location_misalighment_heatmap_reward": 0.7673565149307251,
-      "rewards/overall_score_rewards": 0.8762500286102295,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 96
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 522.0,
-      "epoch": 0.024538325322539842,
-      "grad_norm": 4.630910396575928,
-      "kl": 0.0020904541015625,
-      "learning_rate": 9.8773083733873e-07,
-      "loss": 0.0001,
-      "reward": 5.352054595947266,
-      "reward_std": 0.11839918047189713,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.7040482759475708,
-      "rewards/location_misalighment_heatmap_reward": 0.6780064105987549,
-      "rewards/overall_score_rewards": 0.9700000286102295,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 97
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 493.875,
-      "epoch": 0.02479129774854541,
-      "grad_norm": 3.0994019508361816,
-      "kl": 0.0029296875,
-      "learning_rate": 9.876043511257273e-07,
-      "loss": 0.0001,
-      "reward": 5.197859764099121,
-      "reward_std": 0.10205817222595215,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.6539163589477539,
-      "rewards/location_misalighment_heatmap_reward": 0.6039435267448425,
-      "rewards/overall_score_rewards": 0.940000057220459,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 98
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 495.75,
-      "epoch": 0.025044270174550975,
-      "grad_norm": 2.118175506591797,
-      "kl": 0.0023040771484375,
-      "learning_rate": 9.874778649127244e-07,
-      "loss": 0.0,
-      "reward": 5.037291049957275,
-      "reward_std": 0.07913331687450409,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.4281021058559418,
-      "rewards/location_misalighment_heatmap_reward": 0.7291893362998962,
-      "rewards/overall_score_rewards": 0.8800000548362732,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 99
-    },
-    {
-      "clip_ratio": 0.0,
-      "completion_length": 454.5,
-      "epoch": 0.02529724260055654,
-      "grad_norm": 2.1774909496307373,
-      "kl": 0.0022430419921875,
-      "learning_rate": 9.873513786997217e-07,
-      "loss": -0.0005,
-      "reward": 5.463089466094971,
-      "reward_std": 0.05613892897963524,
-      "rewards/aesthetic_score_rewards": 1.0,
-      "rewards/location_artifact_heatmap_reward": 0.7648959159851074,
-      "rewards/location_misalighment_heatmap_reward": 0.7919434309005737,
-      "rewards/overall_score_rewards": 0.90625,
-      "rewards/plausibility_score_rewards": 1.0,
-      "rewards/semantic_score_rewards": 1.0,
-      "step": 100
-    }
-  ],
-  "logging_steps": 1.0,
-  "max_steps": 7906,
-  "num_input_tokens_seen": 0,
-  "num_train_epochs": 2,
-  "save_steps": 50,
-  "stateful_callbacks": {
-    "TrainerControl": {
-      "args": {
-        "should_epoch_stop": false,
-        "should_evaluate": false,
-        "should_log": false,
-        "should_save": true,
-        "should_training_stop": false
-      },
-      "attributes": {}
-    }
-  },
-  "total_flos": 0.0,
-  "train_batch_size": 1,
-  "trial_name": null,
-  "trial_params": null
-}

else/zero_to_fp32.py

@@ -1,674 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright (c) Microsoft Corporation.
-# SPDX-License-Identifier: Apache-2.0
-
-# DeepSpeed Team
-
-# This script extracts fp32 consolidated weights from a zero 1, 2 and 3 DeepSpeed checkpoints. It gets
-# copied into the top level checkpoint dir, so the user can easily do the conversion at any point in
-# the future. Once extracted, the weights don't require DeepSpeed and can be used in any
-# application.
-#
-# example:
-#   python zero_to_fp32.py . output_dir/
-#   or
-#   python zero_to_fp32.py . output_dir/ --safe_serialization
-
-import argparse
-import torch
-import glob
-import math
-import os
-import re
-import json
-from tqdm import tqdm
-from collections import OrderedDict
-from dataclasses import dataclass
-
-# while this script doesn't use deepspeed to recover data, since the checkpoints are pickled with
-# DeepSpeed data structures it has to be available in the current python environment.
-from deepspeed.utils import logger
-from deepspeed.checkpoint.constants import (DS_VERSION, OPTIMIZER_STATE_DICT, SINGLE_PARTITION_OF_FP32_GROUPS,
-                                            FP32_FLAT_GROUPS, ZERO_STAGE, PARTITION_COUNT, PARAM_SHAPES, BUFFER_NAMES,
-                                            FROZEN_PARAM_SHAPES, FROZEN_PARAM_FRAGMENTS)
-
-
-@dataclass
-class zero_model_state:
-    buffers: dict()
-    param_shapes: dict()
-    shared_params: list
-    ds_version: int
-    frozen_param_shapes: dict()
-    frozen_param_fragments: dict()
-
-
-debug = 0
-
-# load to cpu
-device = torch.device('cpu')
-
-
-def atoi(text):
-    return int(text) if text.isdigit() else text
-
-
-def natural_keys(text):
-    '''
-    alist.sort(key=natural_keys) sorts in human order
-    http://nedbatchelder.com/blog/200712/human_sorting.html
-    (See Toothy's implementation in the comments)
-    '''
-    return [atoi(c) for c in re.split(r'(\d+)', text)]
-
-
-def get_model_state_file(checkpoint_dir, zero_stage):
-    if not os.path.isdir(checkpoint_dir):
-        raise FileNotFoundError(f"Directory '{checkpoint_dir}' doesn't exist")
-
-    # there should be only one file
-    if zero_stage <= 2:
-        file = os.path.join(checkpoint_dir, "mp_rank_00_model_states.pt")
-    elif zero_stage == 3:
-        file = os.path.join(checkpoint_dir, "zero_pp_rank_0_mp_rank_00_model_states.pt")
-
-    if not os.path.exists(file):
-        raise FileNotFoundError(f"can't find model states file at '{file}'")
-
-    return file
-
-
-def get_checkpoint_files(checkpoint_dir, glob_pattern):
-    # XXX: need to test that this simple glob rule works for multi-node setup too
-    ckpt_files = sorted(glob.glob(os.path.join(checkpoint_dir, glob_pattern)), key=natural_keys)
-
-    if len(ckpt_files) == 0:
-        raise FileNotFoundError(f"can't find {glob_pattern} files in directory '{checkpoint_dir}'")
-
-    return ckpt_files
-
-
-def get_optim_files(checkpoint_dir):
-    return get_checkpoint_files(checkpoint_dir, "*_optim_states.pt")
-
-
-def get_model_state_files(checkpoint_dir):
-    return get_checkpoint_files(checkpoint_dir, "*_model_states.pt")
-
-
-def parse_model_states(files):
-    zero_model_states = []
-    for file in files:
-        state_dict = torch.load(file, map_location=device)
-
-        if BUFFER_NAMES not in state_dict:
-            raise ValueError(f"{file} is not a model state checkpoint")
-        buffer_names = state_dict[BUFFER_NAMES]
-        if debug:
-            print("Found buffers:", buffer_names)
-
-        # recover just the buffers while restoring them to fp32 if they were saved in fp16
-        buffers = {k: v.float() for k, v in state_dict["module"].items() if k in buffer_names}
-        param_shapes = state_dict[PARAM_SHAPES]
-
-        # collect parameters that are included in param_shapes
-        param_names = []
-        for s in param_shapes:
-            for name in s.keys():
-                param_names.append(name)
-
-        # update with frozen parameters
-        frozen_param_shapes = state_dict.get(FROZEN_PARAM_SHAPES, None)
-        if frozen_param_shapes is not None:
-            if debug:
-                print(f"Found frozen_param_shapes: {frozen_param_shapes}")
-            param_names += list(frozen_param_shapes.keys())
-
-        # handle shared params
-        shared_params = [[k, v] for k, v in state_dict["shared_params"].items()]
-
-        ds_version = state_dict.get(DS_VERSION, None)
-
-        frozen_param_fragments = state_dict.get(FROZEN_PARAM_FRAGMENTS, None)
-
-        z_model_state = zero_model_state(buffers=buffers,
-                                         param_shapes=param_shapes,
-                                         shared_params=shared_params,
-                                         ds_version=ds_version,
-                                         frozen_param_shapes=frozen_param_shapes,
-                                         frozen_param_fragments=frozen_param_fragments)
-        zero_model_states.append(z_model_state)
-
-    return zero_model_states
-
-
-def parse_optim_states(files, ds_checkpoint_dir):
-    total_files = len(files)
-    state_dicts = []
-    for f in files:
-        state_dict = torch.load(f, map_location=device)
-        # immediately discard the potentially huge 2 optimizer states as we only care for fp32 master weights
-        # and also handle the case where it was already removed by another helper script
-        state_dict["optimizer_state_dict"].pop("optimizer_state_dict", None)
-        state_dicts.append(state_dict)
-
-    if not ZERO_STAGE in state_dicts[0][OPTIMIZER_STATE_DICT]:
-        raise ValueError(f"{files[0]} is not a zero checkpoint")
-    zero_stage = state_dicts[0][OPTIMIZER_STATE_DICT][ZERO_STAGE]
-    world_size = state_dicts[0][OPTIMIZER_STATE_DICT][PARTITION_COUNT]
-
-    # For ZeRO-2 each param group can have different partition_count as data parallelism for expert
-    # parameters can be different from data parallelism for non-expert parameters. So we can just
-    # use the max of the partition_count to get the dp world_size.
-
-    if type(world_size) is list:
-        world_size = max(world_size)
-
-    if world_size != total_files:
-        raise ValueError(
-            f"Expected {world_size} of '*_optim_states.pt' under '{ds_checkpoint_dir}' but found {total_files} files. "
-            "Possibly due to an overwrite of an old checkpoint, or a checkpoint didn't get saved by one or more processes."
-        )
-
-    # the groups are named differently in each stage
-    if zero_stage <= 2:
-        fp32_groups_key = SINGLE_PARTITION_OF_FP32_GROUPS
-    elif zero_stage == 3:
-        fp32_groups_key = FP32_FLAT_GROUPS
-    else:
-        raise ValueError(f"unknown zero stage {zero_stage}")
-
-    if zero_stage <= 2:
-        fp32_flat_groups = [state_dicts[i][OPTIMIZER_STATE_DICT][fp32_groups_key] for i in range(len(state_dicts))]
-    elif zero_stage == 3:
-        # if there is more than one param group, there will be multiple flattened tensors - one
-        # flattened tensor per group - for simplicity merge them into a single tensor
-        #
-        # XXX: could make the script more memory efficient for when there are multiple groups - it
-        # will require matching the sub-lists of param_shapes for each param group flattened tensor
-
-        fp32_flat_groups = [
-            torch.cat(state_dicts[i][OPTIMIZER_STATE_DICT][fp32_groups_key], 0) for i in range(len(state_dicts))
-        ]
-
-    return zero_stage, world_size, fp32_flat_groups
-
-
-def _get_fp32_state_dict_from_zero_checkpoint(ds_checkpoint_dir, exclude_frozen_parameters):
-    """
-    Returns fp32 state_dict reconstructed from ds checkpoint
-
-    Args:
-        - ``ds_checkpoint_dir``: path to the deepspeed checkpoint folder (where the optimizer files are)
-
-    """
-    print(f"Processing zero checkpoint '{ds_checkpoint_dir}'")
-
-    optim_files = get_optim_files(ds_checkpoint_dir)
-    zero_stage, world_size, fp32_flat_groups = parse_optim_states(optim_files, ds_checkpoint_dir)
-    print(f"Detected checkpoint of type zero stage {zero_stage}, world_size: {world_size}")
-
-    model_files = get_model_state_files(ds_checkpoint_dir)
-
-    zero_model_states = parse_model_states(model_files)
-    print(f'Parsing checkpoint created by deepspeed=={zero_model_states[0].ds_version}')
-
-    if zero_stage <= 2:
-        return _get_fp32_state_dict_from_zero2_checkpoint(world_size, fp32_flat_groups, zero_model_states,
-                                                          exclude_frozen_parameters)
-    elif zero_stage == 3:
-        return _get_fp32_state_dict_from_zero3_checkpoint(world_size, fp32_flat_groups, zero_model_states,
-                                                          exclude_frozen_parameters)
-
-
-def _zero2_merge_frozen_params(state_dict, zero_model_states):
-    if zero_model_states[0].frozen_param_shapes is None or len(zero_model_states[0].frozen_param_shapes) == 0:
-        return
-
-    frozen_param_shapes = zero_model_states[0].frozen_param_shapes
-    frozen_param_fragments = zero_model_states[0].frozen_param_fragments
-
-    if debug:
-        num_elem = sum(s.numel() for s in frozen_param_shapes.values())
-        print(f'rank 0: {FROZEN_PARAM_SHAPES}.numel = {num_elem}')
-
-        wanted_params = len(frozen_param_shapes)
-        wanted_numel = sum(s.numel() for s in frozen_param_shapes.values())
-        avail_numel = sum([p.numel() for p in frozen_param_fragments.values()])
-        print(f'Frozen params: Have {avail_numel} numels to process.')
-        print(f'Frozen params: Need {wanted_numel} numels in {wanted_params} params')
-
-    total_params = 0
-    total_numel = 0
-    for name, shape in frozen_param_shapes.items():
-        total_params += 1
-        unpartitioned_numel = shape.numel()
-        total_numel += unpartitioned_numel
-
-        state_dict[name] = frozen_param_fragments[name]
-
-        if debug:
-            print(f"{name} full shape: {shape} unpartitioned numel {unpartitioned_numel} ")
-
-    print(f"Reconstructed Frozen fp32 state dict with {total_params} params {total_numel} elements")
-
-
-def _has_callable(obj, fn):
-    attr = getattr(obj, fn, None)
-    return callable(attr)
-
-
-def _zero2_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states):
-    param_shapes = zero_model_states[0].param_shapes
-
-    # Reconstruction protocol:
-    #
-    # XXX: document this
-
-    if debug:
-        for i in range(world_size):
-            for j in range(len(fp32_flat_groups[0])):
-                print(f"{FP32_FLAT_GROUPS}[{i}][{j}].shape={fp32_flat_groups[i][j].shape}")
-
-    # XXX: memory usage doubles here (zero2)
-    num_param_groups = len(fp32_flat_groups[0])
-    merged_single_partition_of_fp32_groups = []
-    for i in range(num_param_groups):
-        merged_partitions = [sd[i] for sd in fp32_flat_groups]
-        full_single_fp32_vector = torch.cat(merged_partitions, 0)
-        merged_single_partition_of_fp32_groups.append(full_single_fp32_vector)
-    avail_numel = sum(
-        [full_single_fp32_vector.numel() for full_single_fp32_vector in merged_single_partition_of_fp32_groups])
-
-    if debug:
-        wanted_params = sum([len(shapes) for shapes in param_shapes])
-        wanted_numel = sum([sum(shape.numel() for shape in shapes.values()) for shapes in param_shapes])
-        # not asserting if there is a mismatch due to possible padding
-        print(f"Have {avail_numel} numels to process.")
-        print(f"Need {wanted_numel} numels in {wanted_params} params.")
-
-    # params
-    # XXX: for huge models that can't fit into the host's RAM we will have to recode this to support
-    # out-of-core computing solution
-    total_numel = 0
-    total_params = 0
-    for shapes, full_single_fp32_vector in zip(param_shapes, merged_single_partition_of_fp32_groups):
-        offset = 0
-        avail_numel = full_single_fp32_vector.numel()
-        for name, shape in shapes.items():
-
-            unpartitioned_numel = shape.numel() if _has_callable(shape, 'numel') else math.prod(shape)
-            total_numel += unpartitioned_numel
-            total_params += 1
-
-            if debug:
-                print(f"{name} full shape: {shape} unpartitioned numel {unpartitioned_numel} ")
-            state_dict[name] = full_single_fp32_vector.narrow(0, offset, unpartitioned_numel).view(shape)
-            offset += unpartitioned_numel
-
-        # Z2 started to align to 2*world_size to improve nccl performance. Therefore both offset and
-        # avail_numel can differ by anywhere between 0..2*world_size. Due to two unrelated complex
-        # paddings performed in the code it's almost impossible to predict the exact numbers w/o the
-        # live optimizer object, so we are checking that the numbers are within the right range
-        align_to = 2 * world_size
-
-        def zero2_align(x):
-            return align_to * math.ceil(x / align_to)
-
-        if debug:
-            print(f"original offset={offset}, avail_numel={avail_numel}")
-
-        offset = zero2_align(offset)
-        avail_numel = zero2_align(avail_numel)
-
-        if debug:
-            print(f"aligned  offset={offset}, avail_numel={avail_numel}")
-
-        # Sanity check
-        if offset != avail_numel:
-            raise ValueError(f"consumed {offset} numels out of {avail_numel} - something is wrong")
-
-    print(f"Reconstructed fp32 state dict with {total_params} params {total_numel} elements")
-
-
-def _get_fp32_state_dict_from_zero2_checkpoint(world_size, fp32_flat_groups, zero_model_states,
-                                               exclude_frozen_parameters):
-    state_dict = OrderedDict()
-
-    # buffers
-    buffers = zero_model_states[0].buffers
-    state_dict.update(buffers)
-    if debug:
-        print(f"added {len(buffers)} buffers")
-
-    if not exclude_frozen_parameters:
-        _zero2_merge_frozen_params(state_dict, zero_model_states)
-
-    _zero2_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states)
-
-    # recover shared parameters
-    for pair in zero_model_states[0].shared_params:
-        if pair[1] in state_dict:
-            state_dict[pair[0]] = state_dict[pair[1]]
-
-    return state_dict
-
-
-def zero3_partitioned_param_info(unpartitioned_numel, world_size):
-    remainder = unpartitioned_numel % world_size
-    padding_numel = (world_size - remainder) if remainder else 0
-    partitioned_numel = math.ceil(unpartitioned_numel / world_size)
-    return partitioned_numel, padding_numel
-
-
-def _zero3_merge_frozen_params(state_dict, world_size, zero_model_states):
-    if zero_model_states[0].frozen_param_shapes is None or len(zero_model_states[0].frozen_param_shapes) == 0:
-        return
-
-    if debug:
-        for i in range(world_size):
-            num_elem = sum(s.numel() for s in zero_model_states[i].frozen_param_fragments.values())
-            print(f'rank {i}: {FROZEN_PARAM_SHAPES}.numel = {num_elem}')
-
-        frozen_param_shapes = zero_model_states[0].frozen_param_shapes
-        wanted_params = len(frozen_param_shapes)
-        wanted_numel = sum(s.numel() for s in frozen_param_shapes.values())
-        avail_numel = sum([p.numel() for p in zero_model_states[0].frozen_param_fragments.values()]) * world_size
-        print(f'Frozen params: Have {avail_numel} numels to process.')
-        print(f'Frozen params: Need {wanted_numel} numels in {wanted_params} params')
-
-    total_params = 0
-    total_numel = 0
-    for name, shape in zero_model_states[0].frozen_param_shapes.items():
-        total_params += 1
-        unpartitioned_numel = shape.numel()
-        total_numel += unpartitioned_numel
-
-        param_frags = tuple(model_state.frozen_param_fragments[name] for model_state in zero_model_states)
-        state_dict[name] = torch.cat(param_frags, 0).narrow(0, 0, unpartitioned_numel).view(shape)
-
-        partitioned_numel, partitioned_padding_numel = zero3_partitioned_param_info(unpartitioned_numel, world_size)
-
-        if debug:
-            print(
-                f"Frozen params: {total_params} {name} full shape: {shape} partition0 numel={partitioned_numel} partitioned_padding_numel={partitioned_padding_numel}"
-            )
-
-    print(f"Reconstructed Frozen fp32 state dict with {total_params} params {total_numel} elements")
-
-
-def _zero3_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states):
-    param_shapes = zero_model_states[0].param_shapes
-    avail_numel = fp32_flat_groups[0].numel() * world_size
-    # Reconstruction protocol: For zero3 we need to zip the partitions together at boundary of each
-    # param, re-consolidating each param, while dealing with padding if any
-
-    # merge list of dicts, preserving order
-    param_shapes = {k: v for d in param_shapes for k, v in d.items()}
-
-    if debug:
-        for i in range(world_size):
-            print(f"{FP32_FLAT_GROUPS}[{i}].shape={fp32_flat_groups[i].shape}")
-
-        wanted_params = len(param_shapes)
-        wanted_numel = sum(shape.numel() for shape in param_shapes.values())
-        # not asserting if there is a mismatch due to possible padding
-        avail_numel = fp32_flat_groups[0].numel() * world_size
-        print(f"Trainable params: Have {avail_numel} numels to process.")
-        print(f"Trainable params: Need {wanted_numel} numels in {wanted_params} params.")
-
-    # params
-    # XXX: for huge models that can't fit into the host's RAM we will have to recode this to support
-    # out-of-core computing solution
-    offset = 0
-    total_numel = 0
-    total_params = 0
-    for name, shape in tqdm(param_shapes.items(), desc='Gathering Sharded Weights'):
-        unpartitioned_numel = shape.numel()
-        total_numel += unpartitioned_numel
-        total_params += 1
-        partitioned_numel, partitioned_padding_numel = zero3_partitioned_param_info(unpartitioned_numel, world_size)
-
-        if debug:
-            print(
-                f"Trainable params: {total_params} {name} full shape: {shape} partition0 numel={partitioned_numel} partitioned_padding_numel={partitioned_padding_numel}"
-            )
-
-        # XXX: memory usage doubles here
-        state_dict[name] = torch.cat(
-            tuple(fp32_flat_groups[i].narrow(0, offset, partitioned_numel) for i in range(world_size)),
-            0).narrow(0, 0, unpartitioned_numel).view(shape)
-        offset += partitioned_numel
-
-    offset *= world_size
-
-    # Sanity check
-    if offset != avail_numel:
-        raise ValueError(f"consumed {offset} numels out of {avail_numel} - something is wrong")
-
-    print(f"Reconstructed Trainable fp32 state dict with {total_params} params {total_numel} elements")
-
-
-def _get_fp32_state_dict_from_zero3_checkpoint(world_size, fp32_flat_groups, zero_model_states,
-                                               exclude_frozen_parameters):
-    state_dict = OrderedDict()
-
-    # buffers
-    buffers = zero_model_states[0].buffers
-    state_dict.update(buffers)
-    if debug:
-        print(f"added {len(buffers)} buffers")
-
-    if not exclude_frozen_parameters:
-        _zero3_merge_frozen_params(state_dict, world_size, zero_model_states)
-
-    _zero3_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states)
-
-    # recover shared parameters
-    for pair in zero_model_states[0].shared_params:
-        if pair[1] in state_dict:
-            state_dict[pair[0]] = state_dict[pair[1]]
-
-    return state_dict
-
-
-def get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, tag=None, exclude_frozen_parameters=False):
-    """
-    Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated state_dict that can be loaded with
-    ``load_state_dict()`` and used for training without DeepSpeed or shared with others, for example
-    via a model hub.
-
-    Args:
-        - ``checkpoint_dir``: path to the desired checkpoint folder
-        - ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in 'latest' file. e.g., ``global_step14``
-        - ``exclude_frozen_parameters``: exclude frozen parameters
-
-    Returns:
-        - pytorch ``state_dict``
-
-    Note: this approach may not work if your application doesn't have sufficient free CPU memory and
-    you may need to use the offline approach using the ``zero_to_fp32.py`` script that is saved with
-    the checkpoint.
-
-    A typical usage might be ::
-
-        from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint
-        # do the training and checkpoint saving
-        state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir) # already on cpu
-        model = model.cpu() # move to cpu
-        model.load_state_dict(state_dict)
-        # submit to model hub or save the model to share with others
-
-    In this example the ``model`` will no longer be usable in the deepspeed context of the same
-    application. i.e. you will need to re-initialize the deepspeed engine, since
-    ``model.load_state_dict(state_dict)`` will remove all the deepspeed magic from it.
-
-    If you want it all done for you, use ``load_state_dict_from_zero_checkpoint`` instead.
-
-    """
-    if tag is None:
-        latest_path = os.path.join(checkpoint_dir, 'latest')
-        if os.path.isfile(latest_path):
-            with open(latest_path, 'r') as fd:
-                tag = fd.read().strip()
-        else:
-            raise ValueError(f"Unable to find 'latest' file at {latest_path}")
-
-    ds_checkpoint_dir = os.path.join(checkpoint_dir, tag)
-
-    if not os.path.isdir(ds_checkpoint_dir):
-        raise FileNotFoundError(f"Directory '{ds_checkpoint_dir}' doesn't exist")
-
-    return _get_fp32_state_dict_from_zero_checkpoint(ds_checkpoint_dir, exclude_frozen_parameters)
-
-
-def convert_zero_checkpoint_to_fp32_state_dict(checkpoint_dir,
-                                               output_dir,
-                                               max_shard_size="5GB",
-                                               safe_serialization=False,
-                                               tag=None,
-                                               exclude_frozen_parameters=False):
-    """
-    Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated ``state_dict`` file that can be
-    loaded with ``torch.load(file)`` + ``load_state_dict()`` and used for training without DeepSpeed.
-
-    Args:
-        - ``checkpoint_dir``: path to the desired checkpoint folder. (one that contains the tag-folder, like ``global_step14``)
-        - ``output_dir``: directory to the pytorch fp32 state_dict output files
-        - ``max_shard_size``: the maximum size for a checkpoint before being sharded, default value is 5GB
-        - ``safe_serialization``:  whether to save the model using `safetensors` or the traditional PyTorch way (that uses `pickle`).
-        - ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in the file named ``latest`` in the checkpoint folder, e.g., ``global_step14``
-        - ``exclude_frozen_parameters``: exclude frozen parameters
-    """
-    # Dependency pre-check
-    if safe_serialization:
-        try:
-            from safetensors.torch import save_file
-        except ImportError:
-            print('If you want to use `safe_serialization`, please `pip install safetensors`')
-            raise
-    if max_shard_size is not None:
-        try:
-            from huggingface_hub import split_torch_state_dict_into_shards
-        except ImportError:
-            print('If you want to use `max_shard_size`, please `pip install huggingface_hub`')
-            raise
-
-    # Convert zero checkpoint to state_dict
-    state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, tag, exclude_frozen_parameters)
-
-    # Shard the model if it is too big.
-    weights_name = "model.safetensors" if safe_serialization else "pytorch_model.bin"
-    if max_shard_size is not None:
-        filename_pattern = weights_name.replace(".bin", "{suffix}.bin").replace(".safetensors", "{suffix}.safetensors")
-        state_dict_split = split_torch_state_dict_into_shards(state_dict,
-                                                              filename_pattern=filename_pattern,
-                                                              max_shard_size=max_shard_size)
-    else:
-        from collections import namedtuple
-        StateDictSplit = namedtuple("StateDictSplit", ["is_sharded", "filename_to_tensors"])
-        state_dict_split = StateDictSplit(is_sharded=False,
-                                          filename_to_tensors={weights_name: list(state_dict.keys())})
-
-    # Save the model
-    filename_to_tensors = state_dict_split.filename_to_tensors.items()
-    for shard_file, tensors in tqdm(filename_to_tensors, desc="Saving checkpoint shards"):
-        shard = {tensor: state_dict[tensor].contiguous() for tensor in tensors}
-        output_path = os.path.join(output_dir, shard_file)
-        if safe_serialization:
-            save_file(shard, output_path, metadata={"format": "pt"})
-        else:
-            torch.save(shard, output_path)
-
-    # Save index if sharded
-    if state_dict_split.is_sharded:
-        index = {
-            "metadata": state_dict_split.metadata,
-            "weight_map": state_dict_split.tensor_to_filename,
-        }
-        save_index_file = "model.safetensors.index.json" if safe_serialization else "pytorch_model.bin.index.json"
-        save_index_file = os.path.join(output_dir, save_index_file)
-        with open(save_index_file, "w", encoding="utf-8") as f:
-            content = json.dumps(index, indent=2, sort_keys=True) + "\n"
-            f.write(content)
-
-
-def load_state_dict_from_zero_checkpoint(model, checkpoint_dir, tag=None):
-    """
-    1. Put the provided model to cpu
-    2. Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated ``state_dict``
-    3. Load it into the provided model
-
-    Args:
-        - ``model``: the model object to update
-        - ``checkpoint_dir``: path to the desired checkpoint folder. (one that contains the tag-folder, like ``global_step14``)
-        - ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in the file named ``latest`` in the checkpoint folder, e.g., ``global_step14``
-
-    Returns:
-        - ``model`: modified model
-
-    Make sure you have plenty of CPU memory available before you call this function. If you don't
-    have enough use the ``zero_to_fp32.py`` utility to do the conversion. You will find it
-    conveniently placed for you in the checkpoint folder.
-
-    A typical usage might be ::
-
-        from deepspeed.utils.zero_to_fp32 import load_state_dict_from_zero_checkpoint
-        model = load_state_dict_from_zero_checkpoint(trainer.model, checkpoint_dir)
-        # submit to model hub or save the model to share with others
-
-    Note, that once this was run, the ``model`` will no longer be usable in the deepspeed context
-    of the same application. i.e. you will need to re-initialize the deepspeed engine, since
-    ``model.load_state_dict(state_dict)`` will remove all the deepspeed magic from it.
-
-    """
-    logger.info(f"Extracting fp32 weights")
-    state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, tag)
-
-    logger.info(f"Overwriting model with fp32 weights")
-    model = model.cpu()
-    model.load_state_dict(state_dict, strict=False)
-
-    return model
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser()
-    parser.add_argument("checkpoint_dir",
-                        type=str,
-                        help="path to the desired checkpoint folder, e.g., path/checkpoint-12")
-    parser.add_argument("output_dir",
-                        type=str,
-                        help="directory to the pytorch fp32 state_dict output files"
-                        "(e.g. path/checkpoint-12-output/)")
-    parser.add_argument(
-        "--max_shard_size",
-        type=str,
-        default="5GB",
-        help="The maximum size for a checkpoint before being sharded. Checkpoints shard will then be each of size"
-        "lower than this size. If expressed as a string, needs to be digits followed by a unit (like `5MB`"
-        "We default it to 5GB in order for models to be able to run easily on free-tier google colab instances"
-        "without CPU OOM issues.")
-    parser.add_argument(
-        "--safe_serialization",
-        default=False,
-        action='store_true',
-        help="Whether to save the model using `safetensors` or the traditional PyTorch way (that uses `pickle`).")
-    parser.add_argument("-t",
-                        "--tag",
-                        type=str,
-                        default=None,
-                        help="checkpoint tag used as a unique identifier for checkpoint. e.g., global_step1")
-    parser.add_argument("--exclude_frozen_parameters", action='store_true', help="exclude frozen parameters")
-    parser.add_argument("-d", "--debug", action='store_true', help="enable debug")
-    args = parser.parse_args()
-
-    debug = args.debug
-
-    convert_zero_checkpoint_to_fp32_state_dict(args.checkpoint_dir,
-                                               args.output_dir,
-                                               max_shard_size=args.max_shard_size,
-                                               safe_serialization=args.safe_serialization,
-                                               tag=args.tag,
-                                               exclude_frozen_parameters=args.exclude_frozen_parameters)