flux integration test

Author: wwwjn

tests/integration_tests/flux.py

@@ -66,6 +66,7 @@ def run_single_test(test_flavor: OverrideDefinitions, full_path: str, output_dir
     # run_test supports sequence of tests.
     test_name = test_flavor.test_name
     dump_folder_arg = f"--job.dump_folder {output_dir}/{test_name}"
+    custom_job_args = "--job.custom_config_module torchtitan/models/flux/job_config.py"
 
     # Random init encoder for offline testing
     model_arg = "--model.name flux"
@@ -91,8 +92,9 @@ def run_single_test(test_flavor: OverrideDefinitions, full_path: str, output_dir
                 f"scripts/flux_inference/run_infer.sh"
             )
 
-        cmd += " " + model_arg
         cmd += " " + dump_folder_arg
+        cmd += " " + custom_job_args
+        cmd += " " + model_arg
         cmd += " " + random_init_encoder_arg
         cmd += " " + clip_encoder_version_arg
         cmd += " " + t5_encoder_version_arg

torchtitan/models/flux/__init__.py

@@ -8,7 +8,7 @@
 from torchtitan.components.lr_scheduler import build_lr_schedulers
 from torchtitan.components.optimizer import build_optimizers
 
-from torchtitan.datasets.flux_dataset import build_flux_dataloader
+from torchtitan.hf_datasets.flux_dataset import build_flux_dataloader
 from torchtitan.protocols.train_spec import TrainSpec
 from .infra.parallelize import parallelize_flux
 from .model.args import FluxModelArgs