Flux Model

[calvinpelletier]

Context

This adds the main Flux flow-matching model to TorchTune.

NOTE: @pbontrager had mostly finished an implementation of this model before going on leave, so this implementation is temporary to unblock the rest of the Flux PRs. I just copied the code from the official Flux repo with some minimal changes. We can replace it with @pbontrager 's version when he returns.

More Flux PRs coming soon.

Changelog

• Flux flow model implementation and builders
• LoRA builders
• Unit tests
• Removed input image shape assertion from autoencoder (the autoencoder can be used with many image resolutions)
• Some utility functions for predicting the noise in an image latent using the flow model
• Added checkpoint loading/saving logic for the flow model (which is currently a no-op since the temporary torchtune implementation is the same as the huggingface one, see note above)

Usage

tune download black-forest-labs/FLUX.1-dev --output-dir /tmp/flux

from torchtune.models.flux import flux_1_dev_flow_model
from torchtune.training.checkpointing import FullModelHFCheckpointer

model = flux_1_dev_flow_model()
checkpointer = FullModelHFCheckpointer(
    "/tmp/flux",
    ["flux1-dev.safetensors"],
    "FLUX",
    "/tmp/flux_",
)
sd = checkpointer.load_checkpoint()["model"]
model.load_state_dict(sd)

Test plan

Manual testing:

• verified that the torchtune implemention of the flow model has perfect parity with the official implementation and same performance
• verified that the LoRA trains properly

Checklist:

• run pre-commit hooks and linters (make sure you've first installed via pre-commit install)
• add unit tests for any new functionality
• update docstrings for any new or updated methods or classes
• run unit tests via pytest tests
• run recipe tests via pytest tests -m integration_test

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/torchtune/2302

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 3 New Failures, 6 Cancelled Jobs

As of commit 9d90451 with merge base 90fd2d3 ():

NEW FAILURES - The following jobs have failed:

• GPU tests / gpu_test (3.9, stable) (gh)
  Process completed with exit code 2.
• Recipe Tests / recipe_test (3.9) (gh)
  Process completed with exit code 2.
• Unit Test / unit_tests (3.9) (gh)
  Process completed with exit code 2.

CANCELLED JOBS - The following jobs were cancelled. Please retry:

• GPU tests / gpu_test (3.10, stable) (gh)
  ##[error]The operation was canceled.
• GPU tests / gpu_test (3.11, stable) (gh)
  ##[error]The operation was canceled.
• Recipe Tests / recipe_test (3.10) (gh)
  ##[error]The operation was canceled.
• Recipe Tests / recipe_test (3.11) (gh)
• Unit Test / unit_tests (3.10) (gh)
  ##[error]The operation was canceled.
• Unit Test / unit_tests (3.11) (gh)
  ##[error]The operation was canceled.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[ebsmothers]

I'm excited to see this PR! This is just a first pass, most of my comments are around the Flux model components. I don't have a ton to say about the noise prediction yet (I may just need to do more research/see it in action). Still need to look at the single/double stream blocks later, so will probably come back with a few more comments

[ebsmothers]

nit but let's name this _utils.py just for consistency with other similar files

[ebsmothers]

Why do we need to do the QKV unpacking in forward? Seems to me like we should just separate Q, K, and V as part of the checkpoint loading

[ebsmothers]

Separately, I may be missing something.. is there a gap in our existing MultiHeadAttention that necessitates the use of these? I think we should support Q norm, K norm, RoPE, all of that there. But admittedly I'm kind unfamiliar with some of the Flux details

[ebsmothers]

Any reason not to compose this of lower-level modules like we do in TransformerDecoder and elsewhere? E.g.

def __init__(
	self,
	pe_embedder: nn.Module,
	img_in: nn.Module,
	time_in: nn.Module,
	vector_in: nn.Module,
	guidance_in: nn.Module,
	txt_in: nn.Module,
	double_blocks: nn.ModuleList,
	single_blocks: nn.ModuleList,
	final_layer: nn.Module
):

I think one could make a case that the blocks may need some constraints, but I think most other components are single-input, single-output. This would make it a lot easier for someone to swap out e.g. MLPEmbedder with something else.

[ebsmothers]

Should timestep_embedding be parametrized as well?

[ebsmothers]

I think | is only available from Python 3.10 and we still support 3.9 for now. Would just stick with Union for the time being

[ebsmothers]

Same here.. we already have this, right?

[ebsmothers]

Noob q: the only diff between dev and schnell is that dev uses guidance and schnell doesn't?

[ebsmothers]

👀 @felipemello1

[ebsmothers]

I think we've talked about this one for a bit anyways, so great to see it finally materializing. Maybe let's put in modules/peft/_utils.py or something?

[ebsmothers]

One other thing we'll have to watch out for here.. for NF4-quantized layers we also add state dict hooks to prevent higher memory on checkpoint save. It may be worth including that logic in this utility

[ebsmothers]

Not a problem for this PR, but in this big if/elif/.../elif/else I think this (no weight conversion) should actually be the else, not Llama2 (for whatever reason)

[ebsmothers]

Noob q: these are constant in all cases?