Graduate flux from experiment folder to core torchtitan

[wwwjn]

FLUX.1 model is a diffusion model, which is different from language models and needs to extend train.py as needed.

• dataset.py
  • moved to datasets/flux_dataset.py
• tokenizer.py
  • Keep it under models/flux folder
• integration_test.py
  • moved to tests/integration_tests/flux.py
  • Because FLUX used a separate train.py and run_train.sh, I kept a copy of run_tests() instead of generalizing integration_tests/run_test.py
• train.py
  • Keep it under models/flux folder
• validate.py
• sample.py
  • moved to inference/

[tianyu-l]

Thanks! Left some comments.

[tianyu-l]

with this I think you'll need to have the custom job config import in all FLUX OverrideDefinitions. BTW I plan to refactor this so that we don't need to always specify the custom impot.

[wwwjn]

That's a nice catch! I double check our current design - we used the base_config.toml to enable corss-model tests are consistent, but for flux model it has a separate integration test file, so it's ok to directly reuse the debug model, instead of adding bunch of parameters that are not contained in (not only the custom_job_config, but also inference, validation, etc). WDYT?

[tianyu-l]

could you share more context around this change? I vaguely remember someone requested this in an issue but we couldn't reproduce.

[wwwjn]

This a new change need to run FLUX model with HuggingFace encoders. In my recent test, I find the HF encoder parameters are all stride tensor, which is not compatible with FSDP. So I manually change the encoder's parameters to continuous tensor.

[tianyu-l]

could you also rebase onto #1871

[wwwjn]

could you also rebase onto #1871

Sure, will rebase later

[tianyu-l]

Looks really good. Thanks for the refactor! Had some minor comments.

[tianyu-l]

For more logical organization of files, you think we can put clip and t5 config.json into tests/assets/flux_test_encoders/?

[tianyu-l]

need to rebase onto #1851