Upload 3 files

README.md

@@ -119,7 +119,7 @@ This part is actually really easy. You just train it on the normal flow-matching
 
 FLUX models use a text model called T5-XXL to get most of its conditioning for the text-to-image task. Importantly, they pad the text out to either 256 (schnell) or 512 (dev) tokens. 512 tokens is the maximum trained length for the model. By padding, I mean they repeat the last token until the sequence is this length.
 
-This results in the model using these padding tokens to [store information](https://arxiv.org/abs/2309.16588). When you [visualize the attention maps of the tokens in the padding segment of the text encoder](https://github.com/kaibioinfo/FluxAttentionMap/blob/main/attentionmap.ipynb), you can see that about 10-40 tokens shortly after the last token of the text and about 10-40 tokens at the end of the padding contain information which the model uses to make images. Because these are normally used to store information, it means that any prompt long enough to not have some of these padding tokens will end up with degraded performance.
+This results in the model using these padding tokens to [store information](https://arxiv.org/abs/2309.16588). When you [visualize the attention maps of the tokens in the padding segment of the text encoder](https://github.com/kaibioinfo/FluxAttentionMap/resolve/main/attentionmap.ipynb), you can see that about 10-40 tokens shortly after the last token of the text and about 10-40 tokens at the end of the padding contain information which the model uses to make images. Because these are normally used to store information, it means that any prompt long enough to not have some of these padding tokens will end up with degraded performance.
 
 It's easy to prevent this by masking out these padding token during attention. BFL and their engineers know this, but they probably decided against it because it works as is and most fast implementations of attention only work with causal (LLM) types of padding and so would let them train faster.
 
@@ -171,6 +171,8 @@ One problem with diffusion models is that they are [multi-task](https://arxiv.or
 2. Implement multi-rank stratified sampling so that during each step the model trained timesteps were selected per batch based on regions, which normalizes the gradients significantly like using a higher batch size would.
 
 ```py
+from scipy.stats import beta as sp_beta
+
 alpha = 2.0
 beta = 1.6
 num_processes = self.accelerator.num_processes
@@ -192,7 +194,7 @@ No one talks about what datasets they train anymore, but I used open ones from t
 
 ## Training
 
-I started training for over a month on a 5x 3090s and about 500,000 images. I used a 600m LoKr for this. The model looked okay after. Then, I [unexpectedly gained access to 7x H100s for compute resources](https://rundiffusion.com), so I merged my PEFT model in and began training on a new LoKr with 3.2b parameters.
+I started training for over a month on a 5x 3090s and about 500,000 images. I used a [600m LoKr](https://huggingface.co/jimmycarter/LibreFLUX/blob/main/lycoris_config.3090.json) for this at batch size 1 (effective batch size 5 via DDP). The model looked okay after. Then, I [unexpectedly gained access to 7x H100s for compute resources](https://runware.ai), so I merged my PEFT model in and began training on a new LoKr with [3.2b parameters](https://huggingface.co/jimmycarter/LibreFLUX/blob/main/lycoris_config.h100.json). For the 7x H100 run I ran a batch size of 6 (effective batch size 42 via DDP).
 
 ## Post-hoc "EMA"
 
@@ -268,7 +270,7 @@ As far as what I think of the FLUX "open source", many models being trained and
 
 <img src="https://huggingface.co/jimmycarter/LibreFLUX/resolve/main/assets/opensource.png" style="max-width: 100%;">
 
-I would like to thank [RunDiffusion](https://rundiffusion.com) for the H100 access.
+I would like to thank [RunWare](https://runware.ai) for the H100 access.
 
 ## Contacting me and grants
 

lycoris_config.3090.json

@@ -0,0 +1,56 @@
+{
+    "algo": "lokr",
+    "multiplier": 1.0,
+    "linear_dim": 1000000,
+    "linear_alpha": 1,
+    "factor": 2,
+    "full_matrix": true,
+    "apply_preset": {
+        "name_algo_map": {
+            "transformer_blocks.[0-7]*": {
+                "algo": "lokr",
+                "factor": 4,
+                "linear_dim": 1000000,
+                "linear_alpha": 1,
+                "full_matrix": true
+            },
+            "transformer_blocks.[8-15]*": {
+                "algo": "lokr",
+                "factor": 5,
+                "linear_dim": 1000000,
+                "linear_alpha": 1,
+                "full_matrix": true
+            },
+            "transformer_blocks.[16-18]*": {
+                "algo": "lokr",
+                "factor": 10,
+                "linear_dim": 1000000,
+                "linear_alpha": 1,
+                "full_matrix": true
+            },
+            "single_transformer_blocks.[0-15]*": {
+                "algo": "lokr",
+                "factor": 8,
+                "linear_dim": 1000000,
+                "linear_alpha": 1,
+                "full_matrix": true
+            },
+            "single_transformer_blocks.[16-23]*": {
+                "algo": "lokr",
+                "factor": 5,
+                "linear_dim": 1000000,
+                "linear_alpha": 1,
+                "full_matrix": true
+            },
+            "single_transformer_blocks.[24-37]*": {
+                "algo": "lokr",
+                "factor": 4,
+                "linear_dim": 1000000,
+                "linear_alpha": 1,
+                "use_scalar": true,
+                "full_matrix": true
+            }
+        },
+        "use_fnmatch": true
+    }
+}

lycoris_config.h100.json

@@ -0,0 +1,164 @@
+{
+    "algo": "lokr",
+    "multiplier": 1.0,
+    "linear_dim": 1000000,
+    "linear_alpha": 1,
+    "factor": 2,
+    "full_matrix": true,
+    "apply_preset": {
+        "target_module": [
+            "FluxTransformerBlock",
+            "FluxSingleTransformerBlock"
+        ],
+        "name_algo_map": {
+            "transformer_blocks.[0-7].ff*": {
+                "algo": "lokr",
+                "factor": 1,
+                "linear_dim": 1000000,
+                "linear_alpha": 1,
+                "full_matrix": true
+            },
+            "transformer_blocks.[0-7].norm1": {
+                "algo": "lokr",
+                "factor": 1,
+                "linear_dim": 1000000,
+                "linear_alpha": 1,
+                "full_matrix": true
+            },
+            "transformer_blocks.[0-7].norm1_context": {
+                "algo": "lokr",
+                "factor": 1,
+                "linear_dim": 1000000,
+                "linear_alpha": 1,
+                "full_matrix": true
+            },
+            "transformer_blocks.[0-7]*": {
+                "algo": "lokr",
+                "factor": 2,
+                "linear_dim": 1000000,
+                "linear_alpha": 1,
+                "full_matrix": true
+            },
+            "transformer_blocks.[8-15].ff*": {
+                "algo": "lokr",
+                "factor": 1,
+                "linear_dim": 1000000,
+                "linear_alpha": 1,
+                "full_matrix": true
+            },
+            "transformer_blocks.[8-15].norm1": {
+                "algo": "lokr",
+                "factor": 1,
+                "linear_dim": 1000000,
+                "linear_alpha": 1,
+                "full_matrix": true
+            },
+            "transformer_blocks.[8-15].norm1_context": {
+                "algo": "lokr",
+                "factor": 1,
+                "linear_dim": 1000000,
+                "linear_alpha": 1,
+                "full_matrix": true
+            },
+            "transformer_blocks.[8-15]*": {
+                "algo": "lokr",
+                "factor": 2,
+                "linear_dim": 1000000,
+                "linear_alpha": 1,
+                "full_matrix": true
+            },
+            "transformer_blocks.[16-18].ff*": {
+                "algo": "lokr",
+                "factor": 4,
+                "linear_dim": 1000000,
+                "linear_alpha": 1,
+                "full_matrix": true
+            },
+            "transformer_blocks.[16-18].norm1": {
+                "algo": "lokr",
+                "factor": 4,
+                "linear_dim": 1000000,
+                "linear_alpha": 1,
+                "full_matrix": true
+            },
+            "transformer_blocks.[16-18].norm1_context": {
+                "algo": "lokr",
+                "factor": 4,
+                "linear_dim": 1000000,
+                "linear_alpha": 1,
+                "full_matrix": true
+            },
+            "transformer_blocks.[16-18]*": {
+                "algo": "lokr",
+                "factor": 8,
+                "linear_dim": 1000000,
+                "linear_alpha": 1,
+                "full_matrix": true
+            },
+            "single_transformer_blocks.[0-15].ff*": {
+                "algo": "lokr",
+                "factor": 3,
+                "linear_dim": 1000000,
+                "linear_alpha": 1,
+                "full_matrix": true
+            },
+            "single_transformer_blocks.[0-15].norm": {
+                "algo": "lokr",
+                "factor": 3,
+                "linear_dim": 1000000,
+                "linear_alpha": 1,
+                "full_matrix": true
+            },
+            "single_transformer_blocks.[0-15]*": {
+                "algo": "lokr",
+                "factor": 6,
+                "linear_dim": 1000000,
+                "linear_alpha": 1,
+                "full_matrix": true
+            },
+            "single_transformer_blocks.[16-23].ff*": {
+                "algo": "lokr",
+                "factor": 2,
+                "linear_dim": 1000000,
+                "linear_alpha": 1,
+                "full_matrix": true
+            },
+            "single_transformer_blocks.[16-23].norm": {
+                "algo": "lokr",
+                "factor": 2,
+                "linear_dim": 1000000,
+                "linear_alpha": 1,
+                "full_matrix": true
+            },
+            "single_transformer_blocks.[16-23]*": {
+                "algo": "lokr",
+                "factor": 4,
+                "linear_dim": 1000000,
+                "linear_alpha": 1,
+                "full_matrix": true
+            },
+            "single_transformer_blocks.[24-37].ff*": {
+                "algo": "lokr",
+                "factor": 1,
+                "linear_dim": 1000000,
+                "linear_alpha": 1,
+                "full_matrix": true
+            },
+            "single_transformer_blocks.[24-37].norm": {
+                "algo": "lokr",
+                "factor": 1,
+                "linear_dim": 1000000,
+                "linear_alpha": 1,
+                "full_matrix": true
+            },
+            "single_transformer_blocks.[24-37]*": {
+                "algo": "lokr",
+                "factor": 2,
+                "linear_dim": 1000000,
+                "linear_alpha": 1,
+                "full_matrix": true
+            }
+        },
+        "use_fnmatch": true
+    }
+}