diff --git a/examples/flux_lora_train_example01.json b/examples/flux_lora_train_example01.json index af1face..a6a27b6 100644 --- a/examples/flux_lora_train_example01.json +++ b/examples/flux_lora_train_example01.json @@ -5,14 +5,14 @@ { "id": 38, "type": "SetNode", - "pos": { - "0": 1138.6033935546875, - "1": 1.119886875152588 - }, - "size": { - "0": 210, - "1": 58 - }, + "pos": [ + 1138.6033935546875, + 1.119886875152588 + ], + "size": [ + 210, + 58 + ], "flags": { "collapsed": true }, @@ -43,14 +43,14 @@ { "id": 48, "type": "GetNode", - "pos": { - "0": 2630, - "1": 450 - }, - "size": { - "0": 210, - "1": 58 - }, + "pos": [ + 2630, + 450 + ], + "size": [ + 210, + 58 + ], "flags": { "collapsed": true }, @@ -76,14 +76,14 @@ { "id": 61, "type": "PreviewImage", - "pos": { - "0": 3707, - "1": 610 - }, - "size": { - "0": 809.35400390625, - "1": 458.6750793457031 - }, + "pos": [ + 3707, + 610 + ], + "size": [ + 809.35400390625, + 458.6750793457031 + ], "flags": {}, "order": 52, "mode": 0, @@ -97,19 +97,20 @@ "outputs": [], "properties": { "Node name for S&R": "PreviewImage" - } + }, + "widgets_values": [] }, { "id": 63, "type": "GetNode", - "pos": { - "0": 3706.7109375, - "1": 460 - }, - "size": { - "0": 210, - "1": 58 - }, + "pos": [ + 3706.7109375, + 460 + ], + "size": [ + 210, + 58 + ], "flags": { "collapsed": true }, @@ -135,14 +136,14 @@ { "id": 68, "type": "GetNode", - "pos": { - "0": 4765.21875, - "1": 468.45684814453125 - }, - "size": { - "0": 210, - "1": 58 - }, + "pos": [ + 4765.21875, + 468.45684814453125 + ], + "size": [ + 210, + 58 + ], "flags": { "collapsed": true }, @@ -168,14 +169,14 @@ { "id": 70, "type": "VisualizeLoss", - "pos": { - "0": 5586, - "1": -246 - }, - "size": { - "0": 254.40000915527344, - "1": 198 - }, + "pos": [ + 5586, + -246 + ], + "size": [ + 254.40000915527344, + 198 + ], "flags": {}, "order": 58, "mode": 0, @@ -218,14 +219,14 @@ { "id": 73, "type": "Display Any (rgthree)", - "pos": { - "0": 6270, - "1": 660 - }, - "size": { - "0": 210, - "1": 76 - }, + "pos": [ + 6270, + 660 + ], + "size": [ + 210, + 76 + ], "flags": {}, "order": 68, "mode": 2, @@ -248,14 +249,14 @@ { "id": 78, "type": "AddLabel", - "pos": { - "0": 2023, - "1": 1177 - }, - "size": { - "0": 315, - "1": 274 - }, + "pos": [ + 2023, + 1177 + ], + "size": [ + 315, + 274 + ], "flags": { "collapsed": true }, @@ -314,14 +315,14 @@ { "id": 79, "type": "SomethingToString", - "pos": { - "0": 1815, - "1": 1177 - }, - "size": { - "0": 315, - "1": 82 - }, + "pos": [ + 1815, + 1177 + ], + "size": [ + 315, + 82 + ], "flags": { "collapsed": true }, @@ -356,14 +357,14 @@ { "id": 80, "type": "AddLabel", - "pos": { - "0": 2982, - "1": 1177 - }, - "size": { - "0": 315, - "1": 274 - }, + "pos": [ + 2982, + 1177 + ], + "size": [ + 315, + 274 + ], "flags": { "collapsed": true }, @@ -422,14 +423,14 @@ { "id": 81, "type": "SomethingToString", - "pos": { - "0": 2774, - "1": 1177 - }, - "size": { - "0": 315, - "1": 82 - }, + "pos": [ + 2774, + 1177 + ], + "size": [ + 315, + 82 + ], "flags": { "collapsed": true }, @@ -464,14 +465,14 @@ { "id": 82, "type": "SomethingToString", - "pos": { - "0": 3909, - "1": 1177 - }, - "size": { - "0": 315, - "1": 82 - }, + "pos": [ + 3909, + 1177 + ], + "size": [ + 315, + 82 + ], "flags": { "collapsed": true }, @@ -506,14 +507,14 @@ { "id": 83, "type": "AddLabel", - "pos": { - "0": 4130, - "1": 1177 - }, - "size": { - "0": 315, - "1": 274 - }, + "pos": [ + 4130, + 1177 + ], + "size": [ + 315, + 274 + ], "flags": { "collapsed": true }, @@ -572,14 +573,14 @@ { "id": 84, "type": "SomethingToString", - "pos": { - "0": 4963, - "1": 1177 - }, - "size": { - "0": 315, - "1": 82 - }, + "pos": [ + 4963, + 1177 + ], + "size": [ + 315, + 82 + ], "flags": { "collapsed": true }, @@ -614,14 +615,14 @@ { "id": 85, "type": "AddLabel", - "pos": { - "0": 5171, - "1": 1177 - }, - "size": { - "0": 315, - "1": 274 - }, + "pos": [ + 5171, + 1177 + ], + "size": [ + 315, + 274 + ], "flags": { "collapsed": true }, @@ -680,14 +681,14 @@ { "id": 88, "type": "Display Any (rgthree)", - "pos": { - "0": 1143.6033935546875, - "1": 82.11988067626953 - }, - "size": { - "0": 210, - "1": 76 - }, + "pos": [ + 1143.6033935546875, + 82.11988067626953 + ], + "size": [ + 210, + 76 + ], "flags": {}, "order": 26, "mode": 0, @@ -711,14 +712,14 @@ { "id": 89, "type": "UploadToHuggingFace", - "pos": { - "0": 5900, - "1": 660 - }, - "size": { - "0": 315, - "1": 178 - }, + "pos": [ + 5900, + 660 + ], + "size": [ + 315, + 178 + ], "flags": {}, "order": 64, "mode": 2, @@ -767,14 +768,14 @@ { "id": 90, "type": "SaveImage", - "pos": { - "0": 5877, - "1": -60 - }, - "size": { - "0": 574.23046875, - "1": 414.46881103515625 - }, + "pos": [ + 5877, + -60 + ], + "size": [ + 574.23046875, + 414.46881103515625 + ], "flags": {}, "order": 62, "mode": 0, @@ -794,14 +795,14 @@ { "id": 105, "type": "Display Any (rgthree)", - "pos": { - "0": 483, - "1": -811 - }, - "size": { - "0": 1073.7608642578125, - "1": 492.8503112792969 - }, + "pos": [ + 483, + -811 + ], + "size": [ + 1073.7608642578125, + 492.8503112792969 + ], "flags": {}, "order": 27, "mode": 0, @@ -824,14 +825,14 @@ { "id": 109, "type": "TrainDatasetAdd", - "pos": { - "0": -772.3832397460938, - "1": 203.22264099121094 - }, - "size": { - "0": 281.5897521972656, - "1": 338 - }, + "pos": [ + -772.3832397460938, + 203.22264099121094 + ], + "size": [ + 281.5897521972656, + 338 + ], "flags": {}, "order": 20, "mode": 0, @@ -844,7 +845,8 @@ { "name": "regularization", "type": "JSON", - "link": null + "link": null, + "shape": 7 } ], "outputs": [ @@ -877,14 +879,14 @@ { "id": 111, "type": "TrainDatasetAdd", - "pos": { - "0": -472.3832702636719, - "1": 203.22264099121094 - }, - "size": { - "0": 267.5897521972656, - "1": 338 - }, + "pos": [ + -472.3832702636719, + 203.22264099121094 + ], + "size": [ + 267.5897521972656, + 338 + ], "flags": {}, "order": 22, "mode": 0, @@ -897,7 +899,8 @@ { "name": "regularization", "type": "JSON", - "link": null + "link": null, + "shape": 7 } ], "outputs": [ @@ -930,14 +933,14 @@ { "id": 112, "type": "TrainDatasetAdd", - "pos": { - "0": -172.38327026367188, - "1": 203.22264099121094 - }, - "size": { - "0": 259.5897521972656, - "1": 338 - }, + "pos": [ + -172.38327026367188, + 203.22264099121094 + ], + "size": [ + 259.5897521972656, + 338 + ], "flags": {}, "order": 23, "mode": 0, @@ -950,7 +953,8 @@ { "name": "regularization", "type": "JSON", - "link": null + "link": null, + "shape": 7 } ], "outputs": [ @@ -983,14 +987,14 @@ { "id": 113, "type": "Note", - "pos": { - "0": -732, - "1": 63 - }, - "size": { - "0": 462.68292236328125, - "1": 79.98078918457031 - }, + "pos": [ + -732, + 63 + ], + "size": [ + 462.68292236328125, + 79.98078918457031 + ], "flags": {}, "order": 3, "mode": 0, @@ -1008,14 +1012,14 @@ { "id": 115, "type": "Note", - "pos": { - "0": 248.60342407226562, - "1": -89.88011932373047 - }, - "size": { - "0": 462.68292236328125, - "1": 79.98078918457031 - }, + "pos": [ + 248.60342407226562, + -89.88011932373047 + ], + "size": [ + 462.68292236328125, + 79.98078918457031 + ], "flags": {}, "order": 4, "mode": 0, @@ -1033,14 +1037,14 @@ { "id": 117, "type": "ImageConcatFromBatch", - "pos": { - "0": 6690, - "1": 410 - }, - "size": { - "0": 315, - "1": 106 - }, + "pos": [ + 6690, + 410 + ], + "size": [ + 315, + 106 + ], "flags": {}, "order": 21, "mode": 0, @@ -1082,14 +1086,14 @@ { "id": 119, "type": "ImageBatchMulti", - "pos": { - "0": 6820, - "1": 180 - }, - "size": { - "0": 210, - "1": 142 - }, + "pos": [ + 6820, + 180 + ], + "size": [ + 210, + 142 + ], "flags": {}, "order": 18, "mode": 0, @@ -1135,14 +1139,14 @@ { "id": 120, "type": "GetImageSizeAndCount", - "pos": { - "0": 6830, - "1": 120 - }, - "size": { - "0": 210, - "1": 86 - }, + "pos": [ + 6830, + 120 + ], + "size": [ + 210, + 86 + ], "flags": { "collapsed": true }, @@ -1164,19 +1168,19 @@ "shape": 3 }, { - "name": "1024 width", + "name": "width", "type": "INT", "links": null, "shape": 3 }, { - "name": "1072 height", + "name": "height", "type": "INT", "links": null, "shape": 3 }, { - "name": "4 count", + "name": "count", "type": "INT", "links": [ 199 @@ -1187,19 +1191,20 @@ ], "properties": { "Node name for S&R": "GetImageSizeAndCount" - } + }, + "widgets_values": [] }, { "id": 121, "type": "SetNode", - "pos": { - "0": 2170, - "1": 1177 - }, - "size": { - "0": 210, - "1": 58 - }, + "pos": [ + 2170, + 1177 + ], + "size": [ + 210, + 58 + ], "flags": { "collapsed": true }, @@ -1232,14 +1237,14 @@ { "id": 122, "type": "SetNode", - "pos": { - "0": 3128, - "1": 1177 - }, - "size": { - "0": 210, - "1": 58 - }, + "pos": [ + 3128, + 1177 + ], + "size": [ + 210, + 58 + ], "flags": { "collapsed": true }, @@ -1272,14 +1277,14 @@ { "id": 123, "type": "GetNode", - "pos": { - "0": 6640, - "1": 190 - }, - "size": { - "0": 210, - "1": 58 - }, + "pos": [ + 6640, + 190 + ], + "size": [ + 210, + 58 + ], "flags": { "collapsed": true }, @@ -1308,14 +1313,14 @@ { "id": 124, "type": "GetNode", - "pos": { - "0": 6640, - "1": 230 - }, - "size": { - "0": 210, - "1": 58 - }, + "pos": [ + 6640, + 230 + ], + "size": [ + 210, + 58 + ], "flags": { "collapsed": true }, @@ -1343,14 +1348,14 @@ { "id": 125, "type": "SetNode", - "pos": { - "0": 4278, - "1": 1177 - }, - "size": { - "0": 210, - "1": 58 - }, + "pos": [ + 4278, + 1177 + ], + "size": [ + 210, + 58 + ], "flags": { "collapsed": true }, @@ -1384,14 +1389,14 @@ { "id": 126, "type": "GetNode", - "pos": { - "0": 6650, - "1": 280 - }, - "size": { - "0": 210, - "1": 58 - }, + "pos": [ + 6650, + 280 + ], + "size": [ + 210, + 58 + ], "flags": { "collapsed": true }, @@ -1419,14 +1424,14 @@ { "id": 127, "type": "SetNode", - "pos": { - "0": 5319, - "1": 1177 - }, - "size": { - "0": 210, - "1": 58 - }, + "pos": [ + 5319, + 1177 + ], + "size": [ + 210, + 58 + ], "flags": { "collapsed": true }, @@ -1459,14 +1464,14 @@ { "id": 128, "type": "GetNode", - "pos": { - "0": 6640, - "1": 330 - }, - "size": { - "0": 210, - "1": 58 - }, + "pos": [ + 6640, + 330 + ], + "size": [ + 210, + 58 + ], "flags": { "collapsed": true }, @@ -1494,14 +1499,14 @@ { "id": 131, "type": "Note", - "pos": { - "0": 478, - "1": -884 - }, - "size": { - "0": 210, - "1": 58 - }, + "pos": [ + 478, + -884 + ], + "size": [ + 210, + 58 + ], "flags": {}, "order": 9, "mode": 0, @@ -1519,14 +1524,14 @@ { "id": 65, "type": "FluxTrainValidate", - "pos": { - "0": 4775.216796875, - "1": 518.4568481445312 - }, - "size": { - "0": 468.5999755859375, - "1": 46 - }, + "pos": [ + 4775.216796875, + 518.4568481445312 + ], + "size": [ + 468.5999755859375, + 46 + ], "flags": {}, "order": 57, "mode": 0, @@ -1539,7 +1544,8 @@ { "name": "validation_settings", "type": "VALSETTINGS", - "link": 94 + "link": 94, + "shape": 7 } ], "outputs": [ @@ -1566,19 +1572,20 @@ ], "properties": { "Node name for S&R": "FluxTrainValidate" - } + }, + "widgets_values": [] }, { "id": 46, "type": "PreviewImage", - "pos": { - "0": 2654, - "1": 609 - }, - "size": { - "0": 850.0181274414062, - "1": 452.6767578125 - }, + "pos": [ + 2654, + 609 + ], + "size": [ + 850.0181274414062, + 452.6767578125 + ], "flags": {}, "order": 43, "mode": 0, @@ -1592,19 +1599,20 @@ "outputs": [], "properties": { "Node name for S&R": "PreviewImage" - } + }, + "widgets_values": [] }, { "id": 66, "type": "PreviewImage", - "pos": { - "0": 4785, - "1": 628 - }, - "size": { - "0": 850.0181274414062, - "1": 452.6767578125 - }, + "pos": [ + 4785, + 628 + ], + "size": [ + 850.0181274414062, + 452.6767578125 + ], "flags": {}, "order": 60, "mode": 0, @@ -1618,19 +1626,20 @@ "outputs": [], "properties": { "Node name for S&R": "PreviewImage" - } + }, + "widgets_values": [] }, { "id": 37, "type": "FluxTrainValidationSettings", - "pos": { - "0": 775, - "1": 18 - }, - "size": { - "0": 315, - "1": 250 - }, + "pos": [ + 775, + 18 + ], + "size": [ + 315, + 250 + ], "flags": {}, "order": 10, "mode": 0, @@ -1664,14 +1673,14 @@ { "id": 116, "type": "Note", - "pos": { - "0": 776, - "1": -111 - }, - "size": { - "0": 308.08209228515625, - "1": 78.06562805175781 - }, + "pos": [ + 776, + -111 + ], + "size": [ + 308.08209228515625, + 78.06562805175781 + ], "flags": {}, "order": 11, "mode": 0, @@ -1689,14 +1698,14 @@ { "id": 9, "type": "PreviewImage", - "pos": { - "0": 1547, - "1": 596 - }, - "size": { - "0": 891.4732666015625, - "1": 476.6578063964844 - }, + "pos": [ + 1547, + 596 + ], + "size": [ + 891.4732666015625, + 476.6578063964844 + ], "flags": {}, "order": 34, "mode": 0, @@ -1710,19 +1719,20 @@ "outputs": [], "properties": { "Node name for S&R": "PreviewImage" - } + }, + "widgets_values": [] }, { "id": 14, "type": "FluxTrainSave", - "pos": { - "0": 1988, - "1": 256 - }, - "size": { - "0": 393, - "1": 122 - }, + "pos": [ + 1988, + 256 + ], + "size": [ + 393, + 122 + ], "flags": {}, "order": 29, "mode": 0, @@ -1769,14 +1779,14 @@ { "id": 8, "type": "FluxTrainValidate", - "pos": { - "0": 1552, - "1": 500 - }, - "size": { - "0": 468.5999755859375, - "1": 46 - }, + "pos": [ + 1552, + 500 + ], + "size": [ + 468.5999755859375, + 46 + ], "flags": {}, "order": 32, "mode": 0, @@ -1789,7 +1799,8 @@ { "name": "validation_settings", "type": "VALSETTINGS", - "link": 60 + "link": 60, + "shape": 7 } ], "outputs": [ @@ -1815,19 +1826,20 @@ ], "properties": { "Node name for S&R": "FluxTrainValidate" - } + }, + "widgets_values": [] }, { "id": 40, "type": "GetNode", - "pos": { - "0": 1546, - "1": 433 - }, - "size": { - "0": 277.0899353027344, - "1": 58 - }, + "pos": [ + 1546, + 433 + ], + "size": [ + 277.0899353027344, + 58 + ], "flags": { "collapsed": true }, @@ -1853,14 +1865,14 @@ { "id": 45, "type": "FluxTrainValidate", - "pos": { - "0": 2640, - "1": 500 - }, - "size": { - "0": 468.5999755859375, - "1": 46 - }, + "pos": [ + 2640, + 500 + ], + "size": [ + 468.5999755859375, + 46 + ], "flags": {}, "order": 41, "mode": 0, @@ -1873,7 +1885,8 @@ { "name": "validation_settings", "type": "VALSETTINGS", - "link": 69 + "link": 69, + "shape": 7 } ], "outputs": [ @@ -1899,19 +1912,20 @@ ], "properties": { "Node name for S&R": "FluxTrainValidate" - } + }, + "widgets_values": [] }, { "id": 60, "type": "FluxTrainValidate", - "pos": { - "0": 3716.708740234375, - "1": 510 - }, - "size": { - "0": 468.5999755859375, - "1": 46 - }, + "pos": [ + 3716.708740234375, + 510 + ], + "size": [ + 468.5999755859375, + 46 + ], "flags": {}, "order": 50, "mode": 0, @@ -1924,7 +1938,8 @@ { "name": "validation_settings", "type": "VALSETTINGS", - "link": 89 + "link": 89, + "shape": 7 } ], "outputs": [ @@ -1950,19 +1965,20 @@ ], "properties": { "Node name for S&R": "FluxTrainValidate" - } + }, + "widgets_values": [] }, { "id": 47, "type": "FluxTrainSave", - "pos": { - "0": 3114, - "1": 323 - }, - "size": { - "0": 393, - "1": 122 - }, + "pos": [ + 3114, + 323 + ], + "size": [ + 393, + 122 + ], "flags": {}, "order": 37, "mode": 0, @@ -2008,14 +2024,14 @@ { "id": 129, "type": "AddLabel", - "pos": { - "0": 6937, - "1": 60 - }, - "size": { - "0": 315, - "1": 274 - }, + "pos": [ + 6937, + 60 + ], + "size": [ + 315, + 274 + ], "flags": { "collapsed": true }, @@ -2074,14 +2090,14 @@ { "id": 62, "type": "FluxTrainSave", - "pos": { - "0": 4202, - "1": 331 - }, - "size": { - "0": 393, - "1": 122 - }, + "pos": [ + 4202, + 331 + ], + "size": [ + 393, + 122 + ], "flags": {}, "order": 46, "mode": 0, @@ -2127,14 +2143,14 @@ { "id": 134, "type": "FluxTrainSave", - "pos": { - "0": 5275, - "1": 328 - }, - "size": { - "0": 393, - "1": 122 - }, + "pos": [ + 5275, + 328 + ], + "size": [ + 393, + 122 + ], "flags": {}, "order": 54, "mode": 0, @@ -2180,14 +2196,14 @@ { "id": 97, "type": "VisualizeLoss", - "pos": { - "0": 1700, - "1": -650 - }, - "size": { - "0": 303.6300048828125, - "1": 198 - }, + "pos": [ + 1700, + -650 + ], + "size": [ + 303.6300048828125, + 198 + ], "flags": {}, "order": 28, "mode": 0, @@ -2230,14 +2246,14 @@ { "id": 99, "type": "VisualizeLoss", - "pos": { - "0": 2950, - "1": -650 - }, - "size": { - "0": 254.40000915527344, - "1": 198 - }, + "pos": [ + 2950, + -650 + ], + "size": [ + 254.40000915527344, + 198 + ], "flags": {}, "order": 36, "mode": 0, @@ -2280,14 +2296,14 @@ { "id": 101, "type": "VisualizeLoss", - "pos": { - "0": 4090, - "1": -650 - }, - "size": { - "0": 254.40000915527344, - "1": 198 - }, + "pos": [ + 4090, + -650 + ], + "size": [ + 254.40000915527344, + 198 + ], "flags": {}, "order": 45, "mode": 0, @@ -2330,14 +2346,14 @@ { "id": 98, "type": "SaveImage", - "pos": { - "0": 1680, - "1": -340 - }, - "size": { - "0": 645.9608764648438, - "1": 439.37261962890625 - }, + "pos": [ + 1680, + -340 + ], + "size": [ + 645.9608764648438, + 439.37261962890625 + ], "flags": {}, "order": 31, "mode": 0, @@ -2357,14 +2373,14 @@ { "id": 100, "type": "SaveImage", - "pos": { - "0": 2990, - "1": -340 - }, - "size": { - "0": 574.23046875, - "1": 414.46881103515625 - }, + "pos": [ + 2990, + -340 + ], + "size": [ + 574.23046875, + 414.46881103515625 + ], "flags": {}, "order": 40, "mode": 0, @@ -2384,14 +2400,14 @@ { "id": 102, "type": "SaveImage", - "pos": { - "0": 4080, - "1": -340 - }, - "size": { - "0": 574.23046875, - "1": 414.46881103515625 - }, + "pos": [ + 4080, + -340 + ], + "size": [ + 574.23046875, + 414.46881103515625 + ], "flags": {}, "order": 49, "mode": 0, @@ -2411,14 +2427,14 @@ { "id": 95, "type": "OptimizerConfig", - "pos": { - "0": 322, - "1": 385 - }, - "size": { - "0": 315, - "1": 244 - }, + "pos": [ + 322, + 385 + ], + "size": [ + 315, + 244 + ], "flags": {}, "order": 13, "mode": 0, @@ -2450,14 +2466,14 @@ { "id": 74, "type": "Display Any (rgthree)", - "pos": { - "0": 6275, - "1": 492 - }, - "size": { - "0": 358.62896728515625, - "1": 76 - }, + "pos": [ + 6275, + 492 + ], + "size": [ + 358.62896728515625, + 76 + ], "flags": {}, "order": 65, "mode": 0, @@ -2480,14 +2496,14 @@ { "id": 133, "type": "FluxTrainEnd", - "pos": { - "0": 5870, - "1": 492 - }, - "size": { - "0": 317.4000244140625, - "1": 98 - }, + "pos": [ + 5870, + 492 + ], + "size": [ + 317.4000244140625, + 98 + ], "flags": {}, "order": 59, "mode": 0, @@ -2537,14 +2553,14 @@ { "id": 130, "type": "SaveImage", - "pos": { - "0": 7132, - "1": 121 - }, - "size": { - "0": 619.8221435546875, - "1": 714.4110107421875 - }, + "pos": [ + 7132, + 121 + ], + "size": [ + 619.8221435546875, + 714.4110107421875 + ], "flags": {}, "order": 67, "mode": 0, @@ -2564,14 +2580,14 @@ { "id": 64, "type": "FluxTrainLoop", - "pos": { - "0": 4770, - "1": 330 - }, - "size": { - "0": 393, - "1": 78 - }, + "pos": [ + 4770, + 330 + ], + "size": [ + 393, + 78 + ], "flags": {}, "order": 51, "mode": 0, @@ -2614,14 +2630,14 @@ { "id": 59, "type": "FluxTrainLoop", - "pos": { - "0": 3700, - "1": 330 - }, - "size": { - "0": 393, - "1": 78 - }, + "pos": [ + 3700, + 330 + ], + "size": [ + 393, + 78 + ], "flags": {}, "order": 42, "mode": 0, @@ -2665,14 +2681,14 @@ { "id": 44, "type": "FluxTrainLoop", - "pos": { - "0": 2630, - "1": 330 - }, - "size": { - "0": 393, - "1": 78 - }, + "pos": [ + 2630, + 330 + ], + "size": [ + 393, + 78 + ], "flags": {}, "order": 33, "mode": 0, @@ -2716,14 +2732,14 @@ { "id": 4, "type": "FluxTrainLoop", - "pos": { - "0": 1519, - "1": 256 - }, - "size": { - "0": 393, - "1": 78 - }, + "pos": [ + 1519, + 256 + ], + "size": [ + 393, + 78 + ], "flags": {}, "order": 25, "mode": 0, @@ -2767,14 +2783,14 @@ { "id": 135, "type": "StringConstantMultiline", - "pos": { - "0": 319, - "1": 729 - }, - "size": { - "0": 400, - "1": 200 - }, + "pos": [ + 319, + 729 + ], + "size": [ + 400, + 200 + ], "flags": {}, "order": 14, "mode": 0, @@ -2801,14 +2817,14 @@ { "id": 136, "type": "FluxTrainModelSelect", - "pos": { - "0": 251.60342407226562, - "1": 45.11988067626953 - }, - "size": { - "0": 427.607421875, - "1": 137.3937225341797 - }, + "pos": [ + 251.60342407226562, + 45.11988067626953 + ], + "size": [ + 427.607421875, + 137.3937225341797 + ], "flags": {}, "order": 15, "mode": 0, @@ -2846,14 +2862,14 @@ { "id": 137, "type": "TrainDatasetGeneralConfig", - "pos": { - "0": -1122, - "1": 203 - }, - "size": { - "0": 316.3266296386719, - "1": 184.85813903808594 - }, + "pos": [ + -1122, + 203 + ], + "size": [ + 316.3266296386719, + 202 + ], "flags": {}, "order": 16, "mode": 0, @@ -2877,20 +2893,21 @@ false, 0, false, - false + false, + ".txt" ] }, { "id": 107, "type": "InitFluxLoRATraining", - "pos": { - "0": 783, - "1": 326 - }, - "size": { - "0": 477.3700866699219, - "1": 877.820068359375 - }, + "pos": [ + 783, + 326 + ], + "size": [ + 477.3700866699219, + 877.820068359375 + ], "flags": {}, "order": 24, "mode": 0, @@ -2913,12 +2930,14 @@ { "name": "resume_args", "type": "ARGS", - "link": null + "link": null, + "shape": 7 }, { "name": "block_args", "type": "ARGS", - "link": null + "link": null, + "shape": 7 }, { "name": "sample_prompts", @@ -2927,6 +2946,12 @@ "widget": { "name": "sample_prompts" } + }, + { + "name": "loss_args", + "type": "ARGS", + "link": null, + "shape": 7 } ], "outputs": [ @@ -3540,6 +3565,7 @@ ], "groups": [ { + "id": 1, "title": "Train_01", "bounding": [ 1439, @@ -3552,6 +3578,7 @@ "flags": {} }, { + "id": 2, "title": "Settings and init", "bounding": [ 195, @@ -3564,6 +3591,7 @@ "flags": {} }, { + "id": 3, "title": "Train_02", "bounding": [ 2602, @@ -3576,6 +3604,7 @@ "flags": {} }, { + "id": 4, "title": "Train_03", "bounding": [ 3681, @@ -3588,6 +3617,7 @@ "flags": {} }, { + "id": 5, "title": "Train_04", "bounding": [ 4753, @@ -3600,6 +3630,7 @@ "flags": {} }, { + "id": 6, "title": "Dataset", "bounding": [ -1190, @@ -3615,12 +3646,14 @@ "config": {}, "extra": { "ds": { - "scale": 0.6830134553650705, + "scale": 0.9849732675808215, "offset": [ - 1162.5041448258619, - 54.39509897809583 + 97.10777845358007, + -297.2563318885526 ] - } + }, + "VHS_latentpreview": true, + "VHS_latentpreviewrate": 0 }, "version": 0.4 } \ No newline at end of file diff --git a/examples/flux_train_example_01 b/examples/flux_train_example_01 deleted file mode 100644 index a17dfa7..0000000 Binary files a/examples/flux_train_example_01 and /dev/null differ diff --git a/nodes.py b/nodes.py index 4c6a7fc..9f5258d 100644 --- a/nodes.py +++ b/nodes.py @@ -357,6 +357,7 @@ class OptimizerConfigProdigyPlusScheduleFree: "use_stableadamw": ("BOOLEAN",{"default": True, "tooltip": "Scales parameter updates by the root-mean-square of the normalised gradient, in essence identical to Adafactor's gradient scaling. Set to False if the adaptive learning rate never improves."}), "use_cautious" : ("BOOLEAN",{"default": False, "tooltip": "Experimental. Perform 'cautious' updates, as proposed in https://arxiv.org/pdf/2411.16085. Modifies the update to isolate and boost values that align with the current gradient."}), "use_adopt": ("BOOLEAN",{"default": False, "tooltip": "Experimental. Performs a modified step where the second moment is updated after the parameter update, so as not to include the current gradient in the denominator. This is a partial implementation of ADOPT (https://arxiv.org/abs/2411.02853), as we don't have a first moment to use for the update."}), + "use_grams": ("BOOLEAN",{"default": False, "tooltip": "Perform 'grams' updates, as proposed in https://arxiv.org/abs/2412.17107. Modifies the update using sign operations that align with the current gradient. Note that we do not have access to a first moment, so this deviates from the paper (we apply the sign directly to the update). May have a limited effect."}), "stochastic_rounding": ("BOOLEAN",{"default": True, "tooltip": "Use stochastic rounding for bfloat16 weights"}), "extra_optimizer_args": ("STRING",{"multiline": True, "default": "", "tooltip": "additional optimizer args"}), diff --git a/prodigyplusschedulefree/core_optimiser.py b/prodigyplusschedulefree/core_optimiser.py index 6153514..e0aeb95 100644 --- a/prodigyplusschedulefree/core_optimiser.py +++ b/prodigyplusschedulefree/core_optimiser.py @@ -1,23 +1,24 @@ import math import torch -from statistics import mean, harmonic_mean, geometric_mean +from statistics import harmonic_mean class CoreOptimiser(torch.optim.Optimizer): def __init__(self, params, lr=1.0, - betas=(0.9, 0.99), beta3=None, beta4=0, + betas=(0.9, 0.99), beta3=None, weight_decay=0.0, + weight_decay_by_lr=True, use_bias_correction=False, d0=1e-6, d_coef=1.0, prodigy_steps=0, - warmup_steps=0, eps=1e-8, split_groups=True, - split_groups_mean="harmonic_mean", + split_groups_mean=True, factored=True, fused_back_pass=False, use_stableadamw=True, use_muon_pp=False, use_cautious=False, + use_grams=False, use_adopt=False, stochastic_rounding=True): @@ -33,29 +34,38 @@ class CoreOptimiser(torch.optim.Optimizer): raise ValueError("Invalid beta parameter at index 1: {}".format(betas[1])) if beta3 is not None and not 0.0 <= beta3 < 1.0: raise ValueError("Invalid beta3 parameter: {}".format(beta3)) - if beta4 is not None and not 0.0 <= beta4 < 1.0: - raise ValueError("Invalid beta4 parameter: {}".format(beta4)) - if split_groups_mean not in {None, "mean", "harmonic_mean", "geometric_mean"}: - raise ValueError(f"Invalid value for split_groups_mean: '{split_groups_mean}'. Must be one of {None, 'mean', 'harmonic_mean', 'geometric_mean'}") - if use_adopt and use_muon_pp: - print(f"[{self.__class__.__name__}] Muon and ADOPT cannot be used at the same time. Muon has been disabled.") - use_muon_pp = False + self.try_hook_kohya_fbp() - defaults = dict(lr=lr, betas=betas, beta3=beta3, beta4=beta4, + if beta3 is None: + beta3 = betas[1] ** 0.5 + + if eps is None: + print(f"[{self.__class__.__name__}] 'eps' is None, Adam-atan2 enabled.") + if use_stableadamw: + print(f"[{self.__class__.__name__}] 'use_stableadamw' has been disabled (mutually exclusive with Adam-atan2).") + use_stableadamw = False + + if use_cautious and use_grams: + print(f"[{self.__class__.__name__}] 'use_grams' has been disabled (mutually exclusive with 'use_cautious').") + use_grams = False + + defaults = dict(lr=lr, betas=betas, beta3=beta3, eps=eps, weight_decay=weight_decay, - d=d0, d0=d0, d_coef=d_coef, + weight_decay_by_lr=weight_decay_by_lr, + d=d0, d_prev=d0, d0=d0, d_coef=d_coef, k=1, train_mode=True, weight_sum=0, prodigy_steps=prodigy_steps, - warmup_steps=warmup_steps, use_bias_correction=use_bias_correction, d_numerator=0.0, + d_denom=0, factored=factored, use_stableadamw=use_stableadamw, use_muon_pp=use_muon_pp, use_cautious=use_cautious, + use_grams=use_grams, use_adopt=use_adopt, stochastic_rounding=stochastic_rounding) @@ -70,7 +80,7 @@ class CoreOptimiser(torch.optim.Optimizer): self.split_groups_mean = split_groups_mean # Properties for fused backward pass. - self.groups_to_process = None + self.parameters_to_process = None self.shared_d = None self.fused_back_pass = fused_back_pass @@ -110,42 +120,42 @@ class CoreOptimiser(torch.optim.Optimizer): return group['running_d_numerator'], group['running_d_denom'] @torch.no_grad() - def get_d_mean(self, groups, mode): - if mode is None: - return None - elif mode == "harmonic_mean": - return harmonic_mean(group['d'] for group in groups) - elif mode == "geometric_mean": - return geometric_mean(group['d'] for group in groups) - elif mode == "mean": - return mean(group['d'] for group in groups) - - raise ValueError(f"Invalid value for split_groups_mean: '{mode}'. Must be one of {None, 'mean', 'harmonic_mean', 'geometric_mean'}") + def get_d_mean(self): + if self.split_groups and self.split_groups_mean: + return harmonic_mean(group['d'] for group in self.param_groups) + return None + + @torch.no_grad() + def get_d_max(self, group): + if self.split_groups: + return max(group['d'] for group in self.param_groups) + return group['d'] # From: https://github.com/KellerJordan/Muon/blob/master/muon.py @torch.no_grad() def newton_schulz_(self, G, steps=6, eps=1e-7): # Inline reshaping step within the method itself. - original_shape = None - if len(G.shape) > 2: - original_shape = G.shape - G = G.view(G.size(0), -1) + X = G.view(G.size(0), -1) + a, b, c = (3.4445, -4.7750, 2.0315) - X = G.bfloat16() - X /= (X.norm() + eps) # ensure top singular value <= 1 + X = X.to(dtype=torch.bfloat16, copy=True) if G.size(0) > G.size(1): X = X.T + + X /= X.norm().add(eps) # ensure top singular value <= 1 for _ in range(steps): A = X @ X.T B = b * A + c * A @ A X = a * X + B @ X + if G.size(0) > G.size(1): X = X.T - if X is not G: - G.copy_(X) - del X - if original_shape is not None: - G = G.view(*original_shape) + + # Gradient scaling adaptation from: https://github.com/leloykun/adaptive-muon + X = torch.einsum('ij,ij->', G.type_as(X), X).clamp(-1.0, 1.0) * X + G.copy_(X.view_as(G)) + del X + return G # Implementation by Nerogar. From: https://github.com/pytorch/pytorch/issues/120376#issuecomment-1974828905 @@ -167,6 +177,18 @@ class CoreOptimiser(torch.optim.Optimizer): # copy the higher 16 bit into the target tensor target.copy_(result.view(dtype=torch.float32)) + def smart_copy(self, target, source, stochastic_rounding, smart_delete_source): + if target is source: + return + + if stochastic_rounding and target.dtype == torch.bfloat16 and source.dtype == torch.float32: + self.copy_stochastic_(target, source) + else: + target.copy_(source) + + if smart_delete_source: + del source + # Modified Adafactor factorisation implementation by Ross Wightman # https://github.com/huggingface/pytorch-image-models/pull/2320 @torch.no_grad() @@ -192,11 +214,11 @@ class CoreOptimiser(torch.optim.Optimizer): return int(sorted_dims[-2][1]), int(sorted_dims[-1][1]) @torch.no_grad() - def initialise_state(self, p, factored, use_muon_pp): + def initialise_state(self, p, group): raise Exception("Not implemented!") @torch.no_grad() - def initialise_state_internal(self, p, factored, use_muon_pp): + def initialise_state_internal(self, p, group): state = self.state[p] needs_init = len(state) == 0 @@ -206,12 +228,14 @@ class CoreOptimiser(torch.optim.Optimizer): sliced_data = self.get_sliced_tensor(p) # NOTE: We don't initialise z/exp_avg here -- subclass needs to do that. - state['muon'] = use_muon_pp and len(grad.shape) >= 2 and grad.size(0) < 10000 + state['muon'] = group['use_muon_pp'] and len(grad.shape) >= 2 - if not state['muon']: + if state['muon']: + state["rms_sq"] = 0 + else: factored_dims = self.factored_dims( grad.shape, - factored=factored, + factored=group['factored'], min_dim_size_to_factor=32 ) @@ -226,19 +250,19 @@ class CoreOptimiser(torch.optim.Optimizer): # Always store second moment low ranks in fp32 to avoid precision issues. Memory difference # between bf16/fp16 and fp32 is negligible here. state["exp_avg_sq"] = [torch.zeros(row_shape, dtype=torch.float32, device=p.device).detach(), - torch.zeros(col_shape, dtype=torch.float32, device=p.device).detach(), - dr, dc, reduce_dc] + torch.zeros(col_shape, dtype=torch.float32, device=p.device).detach(), + dr, dc, reduce_dc] else: state['exp_avg_sq'] = torch.zeros_like(p, memory_format=torch.preserve_format).detach() - + # If the initial weights are zero, don't bother storing them. - if p.count_nonzero() > 0: + if p.any() > 0: state['p0'] = sliced_data.to(dtype=dtype, memory_format=torch.preserve_format, copy=True).detach() else: state['p0'] = torch.tensor(0.0, dtype=dtype, device=p.device) - + state['s'] = torch.zeros_like(sliced_data, memory_format=torch.preserve_format, dtype=dtype).detach() - + return state, needs_init @torch.no_grad() @@ -249,141 +273,139 @@ class CoreOptimiser(torch.optim.Optimizer): if prodigy_steps > 0 and k >= prodigy_steps: return - beta1, beta2 = group['betas'] - beta3, beta4 = group['beta3'], group['beta4'] - - if beta3 is None: - beta3 = beta2 ** 0.5 - - if beta4 is None: - beta4 = beta1 ** 0.5 - - d = group['d'] - d0 = group['d0'] + d, d0 = group['d'], group['d0'] + d_prev = group['d_prev'] d_coef = group['d_coef'] + beta3 = group['beta3'] running_d_numerator, running_d_denom = self.get_running_values_for_group(group) d_numerator = group['d_numerator'] d_numerator *= beta3 + d_prev = d + d_numerator_item = running_d_numerator.item() d_denom_item = running_d_denom.item() - # Prevent the accumulation of negative values in the numerator in early training. - # We still allow negative updates once progress starts being made, as this is - # important for regulating the adaptive stepsize. - if d_numerator_item > 0 or d > d0: - d_numerator = max(0, d_numerator + d_numerator_item) + # Force Prodigy to be extremely confident before increasing the LR when gradient + # and weights drift. + if d_numerator_item < 0: + if d > d0: + # Prevent the accumulation of negative values in the numerator in early training. + # We still allow negative updates once progress starts being made, as this is + # important for regulating the adaptive stepsize. + d_numerator = min(d_numerator, d_numerator_item) + else: + d_numerator += d_numerator_item + + d_hat = math.atan2(d_coef * d_numerator, d_denom_item) + d = max(d, d_hat) - if d_denom_item > 0: - d_hat = max(math.atan2(d_coef * d_numerator, d_denom_item), d) - d = d * beta4 + d_hat * (1 - beta4) if beta4 > 0 else d_hat - group['d'] = d + group['d_prev'] = d_prev group['d_numerator'] = d_numerator + group['d_denom'] = d_denom_item running_d_numerator.zero_() running_d_denom.zero_() - def on_start_step(self, group): - if self.groups_to_process is None: - # Optimiser hasn't run yet, so initialise. - self.groups_to_process = {i: len(group['params']) for i, group in enumerate(self.param_groups)} - elif len(self.groups_to_process) == 0: - # Start of new optimiser run, so grab updated d. - self.groups_to_process = {i: len(group['params']) for i, group in enumerate(self.param_groups)} + def on_start_step(self): + if self.parameters_to_process is None or self.parameters_to_process == 0: + # Optimiser hasn't run yet (or is starting a new step), so initialise. + self.parameters_to_process = sum(len(group['params']) for group in self.param_groups) + + def on_end_step(self): + self.parameters_to_process -= 1 - if not self.split_groups: - # When groups aren't split, calculate d for the first group, + if self.parameters_to_process == 0: + # Update d for next optimiser step. + if self.split_groups: + i = 0 + for group in self.param_groups: + if group['prodigy_steps'] > 0 and group['k'] == group['prodigy_steps']: + print(f"[{self.__class__.__name__}] Prodigy stepsize adaptation disabled after {group['k']} steps for param_group {i}.") + + self.update_d_and_reset(group) + group['weight_sum'] = group.get('running_weight_sum', 0) + group['k'] += 1 + i += 1 + + self.shared_d = self.get_d_mean() + else: + # When groups aren't split, calculate d for the first group (which collects stats for all groups in non-split mode), # then copy to all other groups. - self.update_d_and_reset(group) - for g in self.param_groups: - g['d'] = group['d'] + first_group = self.param_groups[0] + self.update_d_and_reset(first_group) + + i = 0 + for group in self.param_groups: + if group['prodigy_steps'] > 0 and group['k'] == group['prodigy_steps']: + print(f"[{self.__class__.__name__}] Prodigy stepsize adaptation disabled after {group['k']} steps for param_group {i}.") - self.shared_d = self.get_d_mean(self.param_groups, self.split_groups_mean) if self.split_groups else None + group['d'] = first_group['d'] + group['d_numerator'] = first_group['d_numerator'] + group['d_denom'] = first_group['d_denom'] + group['weight_sum'] = group.get('running_weight_sum', 0) + group['k'] += 1 + i += 1 - def on_end_step(self, group): - group_index = self.param_groups.index(group) - - # Decrement params processed so far. - self.groups_to_process[group_index] -= 1 - - # End of param loop for group, update calculations. - if self.groups_to_process[group_index] == 0: - k = group['k'] - prodigy_steps = group['prodigy_steps'] - if prodigy_steps > 0 and k == prodigy_steps: - print(f"[{self.__class__.__name__}] Prodigy stepsize adaptation disabled after {k} steps for param_group {group_index}.") - - self.groups_to_process.pop(group_index) - if self.split_groups: # When groups are split, calculate per-group d. - self.update_d_and_reset(group) - - group['k'] = k + 1 - return True - - return False def get_dlr(self, group): - lr = group['lr'] - k = group['k'] + return (self.shared_d if self.split_groups and self.shared_d else group['d']) * group['lr'] - warmup_steps = group['warmup_steps'] + def update_prodigy(self, state, group, grad, data, num_scale): + # num_scale is used to compensate the numerator calculations when + # clipping/scaling is applied to the incoming update. If we don't + # do this, it will dampen Prodigy's 'd' predictions. - d = group['d'] - dlr = (self.shared_d if self.split_groups and self.shared_d else d) * lr - - # Apply warmup separate to the denom and numerator updates. - if k < warmup_steps: - dlr *= k / warmup_steps - - return dlr - - def update_prodigy(self, state, group, grad, data, dlr): k = group['k'] prodigy_steps = group['prodigy_steps'] if prodigy_steps <= 0 or k < prodigy_steps: - d, d0 = group['d'], group['d0'] beta3 = group['beta3'] + d, d0 = group['d'], group['d0'] - if beta3 is None: - beta3 = group['betas'][1] ** 0.5 - + # Slow down, rather than speed up, as we approach the + # appropriate LR. + d_k = (d0 / d) * d + sliced_grad = self.get_sliced_tensor(grad) sliced_data = self.get_sliced_tensor(data) running_d_numerator, running_d_denom = self.get_running_values_for_group(group) s = state['s'] + x0_minus = state['p0'] - sliced_data - running_d_numerator.add_(torch.dot(sliced_grad, x0_minus), alpha=(d / d0) * dlr) - del x0_minus - - s.mul_(beta3).add_(sliced_grad, alpha=(d / d0) * dlr) + running_d_numerator.add_(torch.dot(sliced_grad, x0_minus), alpha=d_k * num_scale) + + s.mul_(beta3).add_(sliced_grad, alpha=d_k) running_d_denom.add_(s.abs().sum()) + del x0_minus elif 's' in state: # Free the memory used by Prodigy, as we no longer need it. del state['s'] del state['p0'] - def get_update(self, num, denom, group): - d = group['d'] + def update_(self, num, denom, group): + eps = group['eps'] + + if eps is None: + # Approximate scaling for a regular Adam-style update. + b = self.get_clip_threshold(group) + a = 1 / math.atan(1 / b) - if group['eps'] is None: # Adam-atan2. Use atan2 rather than epsilon and division # for parameter updates (https://arxiv.org/abs/2407.05872). # Has the nice property of "clipping" the gradient as well. - update = num.mul_(d).atan2_(denom) + update = num.atan2_(denom.mul_(b)).mul_(a) else: - # Assume eps as already been added. - update = num.div_(denom).mul_(d) + update = num.div_(denom.add_(eps)) - return update + return update, 1.0 - def get_denom(self, state, group): + def get_denom(self, state): exp_avg_sq = state['exp_avg_sq'] - eps = group['eps'] # Adam EMA updates if isinstance(exp_avg_sq, list): @@ -395,63 +417,104 @@ class CoreOptimiser(torch.optim.Optimizer): denom = row_factor * col_factor else: denom = exp_avg_sq.sqrt() - - if eps is not None: - denom.add_(group['d'] * eps) return denom - - def update_first_moment(self, exp_avg, group, grad): - d = group['d'] + + def update_first_moment(self, state, group, grad): + exp_avg = state['exp_avg'] beta1, _ = group['betas'] - - exp_avg.mul_(beta1).add_(grad, value=d * (1 - beta1)) - return exp_avg - - def update_second_moment(self, state, group, grad, beta2, return_denom=True): - d = group['d'] + + return exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1) + + def update_second_moment(self, state, group, grad, beta2, return_denom=True, denom_before_update=False): exp_avg_sq = state['exp_avg_sq'] - # Adafactor / PaLM beta2 decay. Clip beta2 as per Scaling ViT paper. - if group['use_bias_correction']: - beta2 = min(1 - group['k'] ** -0.8, beta2) + denom = None - one_minus_beta2_d = d * d * (1 - beta2) - + if return_denom and denom_before_update: + denom = self.get_denom(state) + # Adam EMA updates if isinstance(exp_avg_sq, list): row_var, col_var, dr, dc, _ = exp_avg_sq - row_var.mul_(beta2).add_( - grad.norm(dim=dr, keepdim=True).square_().div_(grad.shape[dr]), - alpha=one_minus_beta2_d) - col_var.mul_(beta2).add_( - grad.norm(dim=dc, keepdim=True).square_().div_(grad.shape[dc]), - alpha=one_minus_beta2_d) + row_var.lerp_( + grad.norm(dim=dr, keepdim=True).square_().div_(grad.shape[dr]), + weight=1 - beta2 + ) + col_var.lerp_( + grad.norm(dim=dc, keepdim=True).square_().div_(grad.shape[dc]), + weight=1 - beta2 + ) else: - exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=one_minus_beta2_d) + exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2) - return self.get_denom(state, group) if return_denom else None - - def rms_(self, tensor, rms_min): - if rms_min is not None: - rms = tensor.norm().div(tensor.numel() ** 0.5).add(rms_min) - tensor.div_(rms) - return tensor + if return_denom and denom is None: + denom = self.get_denom(state) - # "Cautious Optimizer (C-Optim): Improving Training with One Line of Code" - # https://github.com/kyleliang919/c-optim - def cautious_(self, update, grad, reuse_grad): - if reuse_grad: - mask = grad.mul_(update) > 0 - else: - mask = grad.mul(update) > 0 + return denom - mask_scale = mask.numel() / mask.sum().add(1) - update.mul_(mask).mul_(mask_scale) - del mask + def get_rms(self, tensor, eps=1e-8): + return tensor.norm().div(tensor.numel() ** 0.5).clamp_min(eps) - return update + def rms_(self, tensor, eps): + return tensor.div_(self.get_rms(tensor, eps)) + + def get_clip_threshold(self, group): + return max(1, 8 * (0.99 ** (group['k'] - 1))) + + def try_hook_kohya_fbp(self): + self.kohya_original_patch_adafactor_fused = None + + try: + # Import and patching will fail if not Kohya. + import library.adafactor_fused + + # Get the original method so we can restore it later. + self.kohya_original_patch_adafactor_fused = library.adafactor_fused.patch_adafactor_fused + + # Define the override. + def prodigy_patch_adafactor_fused(optimizer): + unwrapped_optimiser = None + if hasattr(optimizer, "optimizer"): + # If the optimiser is wrapped, forward the calls to the actual optimiser. + def _step(self, *args, **kwargs): + return self.optimizer.step(*args, **kwargs) + + def _step_param(self, *args, **kwargs): + return self.optimizer.step_param(*args, **kwargs) + + optimizer.step = _step.__get__(optimizer) + optimizer.step_param = _step_param.__get__(optimizer) + unwrapped_optimiser = optimizer.optimizer + else: + unwrapped_optimiser = optimizer + + print(f"[{self.__class__.__name__}] Kohya pipeline detected with fused backward pass. Gradient hook patch successful.") + library.adafactor_fused.patch_adafactor_fused = unwrapped_optimiser.kohya_original_patch_adafactor_fused # Restore the original method. + + unwrapped_optimiser.fused_back_pass = True + unwrapped_optimiser.kohya_original_patch_adafactor_fused = None + + # Patch the method. + library.adafactor_fused.patch_adafactor_fused = prodigy_patch_adafactor_fused + except: + pass + + def try_unhook_kohya_fbp(self): + if self.kohya_original_patch_adafactor_fused is None: + return + + try: + # Import and patching will fail if not Kohya. + import library.adafactor_fused + + # User did not opt for fused backward pass, so remove our hook. + library.adafactor_fused.patch_adafactor_fused = self.kohya_original_patch_adafactor_fused + except: + pass + + self.kohya_original_patch_adafactor_fused = None @torch.no_grad() def step_param(self, p, group): @@ -463,6 +526,8 @@ class CoreOptimiser(torch.optim.Optimizer): @torch.no_grad() def step(self, closure=None): + self.try_unhook_kohya_fbp() + if self.fused_back_pass: return diff --git a/prodigyplusschedulefree/prodigy_plus_schedulefree.py b/prodigyplusschedulefree/prodigy_plus_schedulefree.py index 5a26984..7b384b5 100644 --- a/prodigyplusschedulefree/prodigy_plus_schedulefree.py +++ b/prodigyplusschedulefree/prodigy_plus_schedulefree.py @@ -5,8 +5,7 @@ from .core_optimiser import CoreOptimiser class ProdigyPlusScheduleFree(CoreOptimiser): r""" An optimiser based on Prodigy that includes schedule-free logic. Has additional improvements in the form of optional StableAdamW - gradient scaling and Adam-atan2 updates, per parameter group adaptation, lower memory utilisation, fused back pass support and - tweaks to mitigate uncontrolled LR growth. + gradient scaling and Adam-atan2 updates, per parameter group adaptation, lower memory utilisation and fused back pass support. Based on code from: https://github.com/facebookresearch/schedule_free @@ -26,120 +25,132 @@ class ProdigyPlusScheduleFree(CoreOptimiser): ability for the optimiser to predict stepsizes. Gradient clipping/normalisation is already handled in the following configurations: 1) `use_stableadamw=True,eps=1e8` (or any reasonable positive epsilon) - 2) `eps=None` (Adam-atan2, scale invariant and can mess with Prodigy's stepsize calculations in some scenarios) + 2) `eps=None` (Adam-atan2, scale invariant. Will disable StableAdamW if enabled.) - A new parameter, `beta4`, allows `d` to be updated via a moving average, rather than being immediately updated. This can help - smooth out learning rate adjustments. Values of 0.9-0.99 are recommended if trying out the feature. If set to None, the - square root of `beta1` is used, while a setting of 0 (the default) disables the feature. - - By default, `split_groups` is set to `True`, so each parameter group will have its own adaptation values. So if you're training - different networks together, they won't contaminate each other's learning rates. The disadvantage of this approach is that some - networks can take a long time to reach a good learning rate when trained alongside others (for example, SDXL's Unet). - It's recommended to use a higher `d0` (1e-5, 5e-5, 1e-4) so these networks don't get stuck at a low learning rate. + By default, `split_groups` and `split_groups_mean` are set to `True`, so each parameter group will have its own `d` values, however, + they will all use the harmonic mean for the dynamic learning rate. To make each group use its own dynamic LR, set `split_groups_mean` to False. + To use the reference Prodigy behaviour where all groups are combined, set `split_groups` to False. - For Prodigy's reference behaviour, which lumps all parameter groups together, set `split_groups` to `False`. - In some scenarios, it can be advantageous to freeze Prodigy's adaptive stepsize after a certain number of steps. This - can be controlled via the `prodigy_steps` settings. + can be controlled via the `prodigy_steps` settings. This will also free any Prodigy-specific memory used by the + optimiser (though with all the memory-related improvements, this should not be significant unless you're training + very large models). Arguments: params (iterable): Iterable of parameters to optimize or dicts defining parameter groups. lr (float): Learning rate adjustment parameter. Increases or decreases the Prodigy learning rate. + (default: 1.0) betas (Tuple[float, float], optional): - Coefficients used for computing running averages of gradient and its square + Coefficients used for computing running averages of gradient and its square. (default: (0.9, 0.99)) eps (float): Term added to the denominator outside of the root operation to improve numerical stability. If set to None, Adam-atan2 is used instead. This removes the need for epsilon tuning, but may not work well in all situations. (default: 1e-8). beta3 (float): - Coefficient for computing the Prodigy stepsize using running averages. - If set to None, uses the value of square root of beta2 (default: None). - beta4 (float): - Coefficient for updating the learning rate from Prodigy's adaptive stepsize. Smooths out spikes in learning rate adjustments. - If set to None, beta1 is used instead. (default 0, which disables smoothing and uses original Prodigy behaviour). + Coefficient for computing the Prodigy stepsize using running averages. If set to None, uses the value of + square root of beta2 + (default: None). weight_decay (float): - Decoupled weight decay. Value is multiplied by the adaptive learning rate. + Decoupled weight decay. Use the weight_decay_by_lr setting to determine if decay should be multiplied by the + adaptive learning rate. (default: 0). + weight_decay_by_lr (boolean): + If True, weight_decay is multiplied by the adaptive learning rate (as per the PyTorch implementation of AdamW). + If False, weight_decay will have a much stronger effect. + (default: True). use_bias_correction (boolean): - Turn on Adafactor-style bias correction, which scales beta2 directly. (default False). + Turn on Adafactor-style bias correction, which scales beta2 directly. (default: False). d0 (float): - Initial estimate for Prodigy (default 1e-6). + Initial estimate for Prodigy. Also serves as the minimum learning rate. + (default: 1e-6). d_coef (float): - Coefficient in the expression for the estimate of d (default 1.0). Values such as 0.5 and 2.0 typically work as well. + Coefficient in the expression for the estimate of d. Values such as 0.5 and 2.0 typically work as well. Changing this parameter is the preferred way to tune the method. + (default: 1.0) prodigy_steps (int): - Freeze Prodigy stepsize adjustments after a certain optimiser step. - (default 0) - warmup_steps (int): - Enables a linear learning rate warmup (default 0). Use this over the warmup settings of your LR scheduler. + Freeze Prodigy stepsize adjustments after a certain optimiser step and releases all state memory required + by Prodigy. + (default: 0) split_groups (boolean): Track individual adaptation values for each parameter group. For example, if training a text encoder beside a Unet. Note this can have a significant impact on training dynamics. Set to False for original Prodigy behaviour, where all groups share the same values. - (default True) - split_groups_mean (str: None, "mean", "harmonic_mean", "geometric_mean"): - When split_groups is True, use specified mean of learning rates for all groups. This favours + (default: True) + split_groups_mean (boolean): + When split_groups is True, use the harmonic mean of learning rates for all groups. This favours a more conservative LR. Calculation remains per-group. If split_groups is False, this value has no effect. - Set to None to have each group use its own learning rate calculation. - (default "harmonic_mean") + Set to False to have each group use its own learning rate. + (default: True) factored (boolean): Use factored approximation of the second moment, similar to Adafactor. Reduces memory usage. Disable if training results in NaNs or the learning rate fails to grow. - (default True) + (default: True) fused_back_pass (boolean): - Stops the optimiser from running the normal step method. Set to True if using fused backward pass. - (default False) + Stops the optimiser from running the normal step method. Set to True if using fused backward pass. Really only + needed for scripts and UIs that call the regular step method even when using fused backward pass (OneTrainer). + (default: False) use_stableadamw (boolean): Scales parameter updates by the root-mean-square of the normalised gradient, in essence identical to Adafactor's gradient scaling. Set to False if the adaptive learning rate never improves. - (default True) + (default: True) use_muon_pp (boolean): - Experimental. Perform orthogonalisation post-processing on 2D+ parameter updates ala Shampoo/SOAP/Muon. + Experimental. Perform orthogonalisation on the gradient before it is used for updates ala Shampoo/SOAP/Muon. (https://github.com/KellerJordan/Muon/blob/master/muon.py). Not suitable for all training scenarios. - May not work well with small batch sizes or finetuning. (default False) + May not work well with small batch sizes or finetuning. + (default: False) use_cautious (boolean): Experimental. Perform "cautious" updates, as proposed in https://arxiv.org/pdf/2411.16085. Modifies - the update to isolate and boost values that align with the current gradient. - (default False) + the update to isolate and boost values that align with the current gradient. Note that we do not have + access to a first moment, so this deviates from the paper (we apply the mask directly to the update). + May have a limited effect. + (default: False) + use_grams (boolean): + Experimental. Perform "grams" updates, as proposed in https://arxiv.org/abs/2412.17107. Modifies + the update using sign operations that align with the current gradient. Note that we do not have + access to a first moment, so this deviates from the paper (we apply the sign directly to the update). + May have a limited effect. + (default: False) use_adopt (boolean): Experimental. Performs a modified step where the second moment is updated after the parameter update, so as not to include the current gradient in the denominator. This is a partial implementation of ADOPT (https://arxiv.org/abs/2411.02853), as we don't have a first moment to use for the update. - (default False) + (default: False) stochastic_rounding (boolean): Use stochastic rounding for bfloat16 weights (https://github.com/pytorch/pytorch/issues/120376). Brings bfloat16 training performance close to that of float32. - (default True) + (default: True) """ def __init__(self, params, lr=1.0, - betas=(0.9, 0.99), beta3=None, beta4=0, + betas=(0.9, 0.99), beta3=None, weight_decay=0.0, + weight_decay_by_lr=True, use_bias_correction=False, d0=1e-6, d_coef=1.0, prodigy_steps=0, - warmup_steps=0, eps=1e-8, split_groups=True, - split_groups_mean="harmonic_mean", + split_groups_mean=True, factored=True, fused_back_pass=False, use_stableadamw=True, use_muon_pp=False, use_cautious=False, + use_grams=False, use_adopt=False, stochastic_rounding=True): - super().__init__(params=params, lr=lr, betas=betas, beta3=beta3, beta4=beta4, - weight_decay=weight_decay, use_bias_correction=use_bias_correction, - d0=d0, d_coef=d_coef, prodigy_steps=prodigy_steps, - warmup_steps=warmup_steps, eps=eps, split_groups=split_groups, - split_groups_mean=split_groups_mean, factored=factored, - fused_back_pass=fused_back_pass, use_stableadamw=use_stableadamw, - use_muon_pp=use_muon_pp, use_cautious=use_cautious, use_adopt=use_adopt, - stochastic_rounding=stochastic_rounding) + super().__init__(params=params, lr=lr, betas=betas, beta3=beta3, + weight_decay=weight_decay, weight_decay_by_lr=weight_decay_by_lr, + use_bias_correction=use_bias_correction, + d0=d0, d_coef=d_coef, prodigy_steps=prodigy_steps, + eps=eps, split_groups=split_groups, + split_groups_mean=split_groups_mean, factored=factored, + fused_back_pass=fused_back_pass, use_stableadamw=use_stableadamw, + use_muon_pp=use_muon_pp, use_cautious=use_cautious, use_grams=use_grams, + use_adopt=use_adopt, stochastic_rounding=stochastic_rounding) @torch.no_grad() def eval(self): @@ -168,96 +179,116 @@ class ProdigyPlusScheduleFree(CoreOptimiser): group['train_mode'] = True @torch.no_grad() - def initialise_state(self, p, factored, use_muon_pp): - state, needs_init = self.initialise_state_internal(p, factored, use_muon_pp) + def initialise_state(self, p, group): + state, needs_init = self.initialise_state_internal(p, group) if needs_init: - state['z'] = p.detach().clone(memory_format=torch.preserve_format) + state['z'] = p.detach().clone(memory_format=torch.preserve_format) return state - - @torch.no_grad() - def update_params(self, y, z, update, dlr, group): - # Weight decay. - weight_decay = group['weight_decay'] - if weight_decay != 0: - update.add_(y, alpha=weight_decay) + @torch.no_grad() + def update_params(self, y, z, update, group): + dlr = self.get_dlr(group) + + beta1, _ = group['betas'] + decay = group['weight_decay'] weight = dlr ** 2 weight_sum = group['weight_sum'] + weight ckp1 = weight / weight_sum if weight_sum else 0 - y.lerp_(end=z, weight=ckp1) - y.add_(update, alpha=dlr * (group['betas'][0] * (1 - ckp1) - 1)) - z.sub_(update, alpha=dlr) + xy_step = 1 - beta1 * (1 - ckp1) + + if decay != 0: + # Weight decay at Y. + if group['weight_decay_by_lr']: + decay *= dlr + + y.sub_(y, alpha=decay * xy_step) + z.sub_(y, alpha=decay) + + if group['use_cautious']: + # "Cautious Optimizer (C-Optim): Improving Training with One Line of Code": https://github.com/kyleliang919/c-optim + # ScheduleFree implementation by nhamanasu: https://github.com/facebookresearch/schedule_free/pull/54 + u = (y - z).mul_(ckp1).add_(update, alpha=dlr * xy_step) + z.sub_(update, alpha=dlr) + mask = (u * update > 0).to(update.dtype) + mask.mul_(mask.numel() / (mask.sum() + 1)) + u.mul_(mask) + y.sub_(u) + del mask, u + elif group['use_grams']: + # "Grams: Gradient Descent with Adaptive Momentum Scaling": https://arxiv.org/abs/2412.17107 + u = (y - z).mul_(ckp1).add_(update, alpha=dlr * xy_step) + z.sub_(update, alpha=dlr) # Update z now so we can do sign in-place. + y.sub_(u.abs_().mul_(update.sign_())) + del u + else: + y.lerp_(end=z, weight=ckp1) + y.sub_(update, alpha=dlr * xy_step) + z.sub_(update, alpha=dlr) return weight_sum @torch.no_grad() def step_param(self, p, group): + self.on_start_step() + if not group['train_mode']: raise Exception("Not in train mode!") - self.on_start_step(group) - weight_sum = group['weight_sum'] - + if p.grad is not None: - grad = p.grad + grad = p.grad.to(dtype=torch.float32, copy=True) - state = self.initialise_state(p, group['factored'], group['use_muon_pp']) - use_adopt = group['use_adopt'] + use_adopt = group['use_adopt'] + stochastic = group['stochastic_rounding'] + _, beta2 = group['betas'] + k = group['k'] - if use_adopt and group['k'] == 1: - self.update_second_moment(state, group, grad.float(), 0, return_denom=False) + state = self.initialise_state(p, group) + + update = None + + if state['muon']: + grad = self.newton_schulz_(grad) + grad_rms = self.get_rms(grad).item() ** 2 + + rms_sq = (state["rms_sq"] * beta2) + (grad_rms * (1 - beta2)) + state["rms_sq"] = rms_sq + + update = grad.mul_(1.0 / ((rms_sq ** 0.5) + 1e-12)) else: - dlr = self.get_dlr(group) - rms_min = 1.0 if group['use_stableadamw'] else None - y, z = p, state['z'] + if group['use_bias_correction']: + # Adafactor / PaLM beta2 decay. Clip beta2 as per Scaling ViT paper. + beta2 = min(beta2, 1 - k ** -0.8) + beta2 = (1 - beta2) / (1 - beta2 ** k) - self.update_prodigy(state, group, grad, z, dlr) - - grad_mask = grad.clone() if group['use_cautious'] else None - - if state['muon']: - # newton_schulz_ casts to bf16 internally, so do float cast afterwards. - update = self.newton_schulz_(grad).float() - rms_min = 1e-30 + if use_adopt and group['k'] == 1: + self.update_second_moment(state, group, grad, 0, return_denom=False) else: - grad = grad.float() - _, beta2 = group['betas'] - - if use_adopt: - denom = self.get_denom(state, group) - self.update_second_moment(state, group, grad, beta2, return_denom=False) - else: - denom = self.update_second_moment(state, group, grad, beta2) - - update = self.get_update(grad, denom, group) + denom = self.update_second_moment(state, group, grad, beta2, denom_before_update=use_adopt) + update, num_scale = self.update_(grad, denom, group) del denom - if group['eps'] is None: - rms_min = None + if update is not None: + if group['use_stableadamw']: + clip_threshold = self.get_clip_threshold(group) + num_scale = max(1, self.get_rms(update, 1.0).item() / clip_threshold) + update.mul_(1 / num_scale) - self.rms_(update, rms_min) + z_state = state['z'] + self.update_prodigy(state, group, p.grad, z_state, 1.0) - if grad_mask is not None: - self.cautious_(update, grad_mask, reuse_grad=True) + y, z = (p.float(), z_state.float()) if stochastic else (p, z_state) + weight_sum = self.update_params(y, z, update, group) - if group['stochastic_rounding'] and y.dtype == z.dtype == torch.bfloat16: - y_fp32, z_fp32 = y.float(), z.float() - - weight_sum = self.update_params(y_fp32, z_fp32, update, dlr, group) - - self.copy_stochastic_(y, y_fp32) - self.copy_stochastic_(z, z_fp32) - - del y_fp32, z_fp32 - else: - weight_sum = self.update_params(y, z, update, dlr, group) + self.smart_copy(p, y, stochastic, True) + self.smart_copy(z_state, z, stochastic, True) del update - if self.on_end_step(group): - group['weight_sum'] = weight_sum \ No newline at end of file + group['running_weight_sum'] = weight_sum + self.on_end_step() \ No newline at end of file