Update FLUX model configuration and recipes to use 12B variant

Signed-off-by: Ao Tang <aot@nvidia.com>

Author: suiyoubi

examples/diffusion/recipes/flux/README.md

@@ -92,23 +92,23 @@ uv run python examples/diffusion/recipes/flux/inference_flux.py \
 
 Run FLUX pretraining with the generic **run_recipe** script (same entry point as for LLM training).
 
-**Recipe:** [megatron.bridge.diffusion.recipes.flux.flux.pretrain_config](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/src/megatron/bridge/diffusion/recipes/flux/flux.py)
+**Recipe:** [megatron.bridge.diffusion.recipes.flux.flux.flux_12b_pretrain_config](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/src/megatron/bridge/diffusion/recipes/flux/flux.py)
 
 From the **Megatron-Bridge repository root**:
 
 **Mock data (no dataset path):**
 
 ```bash
 uv run torchrun --nproc_per_node=8 scripts/training/run_recipe.py \
-  --recipe flux_14b_pretrain_config \
+  --recipe flux_12b_pretrain_config \
   --step_func flux_step
 ```
 
 **Real data (WebDataset path):** Set `dataset.path` so the recipe uses real data instead of mock:
 
 ```bash
 uv run torchrun --nproc_per_node=8 scripts/training/run_recipe.py \
-  --recipe flux_14b_pretrain_config \
+  --recipe flux_12b_pretrain_config \
   --step_func flux_step \
   dataset.path=${WORKSPACE}/data/my_flux_wds/
 ```
@@ -117,7 +117,7 @@ uv run torchrun --nproc_per_node=8 scripts/training/run_recipe.py \
 
 ```bash
 uv run torchrun --nproc_per_node=8 scripts/training/run_recipe.py \
-  --recipe flux_14b_pretrain_config \
+  --recipe flux_12b_pretrain_config \
   --step_func flux_step \
   dataset.path=${WORKSPACE}/data/my_flux_wds/ \
   train.train_iters=10000 \
@@ -139,7 +139,7 @@ Run FLUX fine-tuning with the generic **run_recipe** script. Set the pretrained
 
 ```bash
 uv run torchrun --nproc_per_node=8 scripts/training/run_recipe.py \
-  --recipe flux_finetune_config \
+  --recipe flux_12b_finetune_config \
   --step_func flux_step \
   checkpoint.pretrained_checkpoint=${WORKSPACE}/checkpoints/flux/flux.1-dev/iter_0000000
 ```
@@ -148,7 +148,7 @@ uv run torchrun --nproc_per_node=8 scripts/training/run_recipe.py \
 
 ```bash
 uv run torchrun --nproc_per_node=8 scripts/training/run_recipe.py \
-  --recipe flux_finetune_config \
+  --recipe flux_12b_finetune_config \
   --step_func flux_step \
   checkpoint.pretrained_checkpoint=${WORKSPACE}/checkpoints/flux/flux.1-dev/iter_0000000 \
   dataset.path=${WORKSPACE}/data/my_flux_wds/ \
@@ -169,9 +169,9 @@ uv run torchrun --nproc_per_node=8 scripts/training/run_recipe.py \
    Run [inference_flux.py](inference_flux.py) with `--flux_ckpt` (Megatron `iter_*` path), `--vae_ckpt`, and `--prompts`.
 
 3. **Pretraining**  
-   Run `scripts/training/run_recipe.py --recipe flux_14b_pretrain_config --step_func flux_step` (optionally with `dataset.path=...` for real data and CLI overrides).
+   Run `scripts/training/run_recipe.py --recipe flux_12b_pretrain_config --step_func flux_step` (optionally with `dataset.path=...` for real data and CLI overrides).
 
 4. **Fine-Tuning**  
-   Run `scripts/training/run_recipe.py --recipe flux_finetune_config --step_func flux_step checkpoint.pretrained_checkpoint=<path>` (optionally with `dataset.path=...` and overrides).
+   Run `scripts/training/run_recipe.py --recipe flux_12b_finetune_config --step_func flux_step checkpoint.pretrained_checkpoint=<path>` (optionally with `dataset.path=...` and overrides).
 
 For more details, see the recipe in `src/megatron/bridge/diffusion/recipes/flux/flux.py` and `scripts/training/run_recipe.py`.

scripts/training/run_recipe.py

@@ -31,7 +31,7 @@
 
     Diffusion (FLUX) pretrain:
         torchrun --nproc_per_node=8 run_recipe.py \
-            --recipe flux_14b_pretrain_config \
+            --recipe flux_12b_pretrain_config \
             --step_func flux_step
 
     With CLI overrides:

src/megatron/bridge/diffusion/data/flux/flux_energon_datamodule.py

@@ -16,7 +16,7 @@
 
 import logging
 from dataclasses import dataclass
-from typing import Optional, Union
+from typing import Optional
 
 from megatron.bridge.data.utils import DatasetBuildContext, DatasetProvider
 from megatron.bridge.diffusion.data.common.diffusion_energon_datamodule import (
@@ -29,17 +29,6 @@
 logger = logging.getLogger(__name__)
 
 
-def _is_empty_path(path: Optional[Union[str, list]]) -> bool:
-    """Return True if path should be treated as empty (use mock data)."""
-    if path is None:
-        return True
-    if isinstance(path, str):
-        return not path.strip()
-    if isinstance(path, list):
-        return len(path) == 0
-    return True
-
-
 @dataclass(kw_only=True)
 class FluxDatasetConfig(DatasetProvider):
     """
@@ -50,7 +39,7 @@ class FluxDatasetConfig(DatasetProvider):
     to load real data; no separate --data_paths or mock flag needed.
     """
 
-    path: Optional[Union[str, list]] = None
+    path: Optional[str] = None
     seq_length: int = 1024
     packing_buffer_size: Optional[int] = None
     micro_batch_size: int = 1
@@ -70,7 +59,7 @@ def __post_init__(self):
         self.sequence_length = self.seq_length
 
     def build_datasets(self, context: DatasetBuildContext):
-        if _is_empty_path(self.path):
+        if not (self.path or "").strip():
             logger.info(
                 "FLUX dataset: path is None or empty; using mock/synthetic data. "
                 "Set dataset.path=/path/to/wds to use real data."
@@ -97,7 +86,6 @@ def build_datasets(self, context: DatasetBuildContext):
             )
             return mock_cfg.build_datasets(context)
 
-        # Real data: path is set (string or list)
         real_cfg = FluxDataModuleConfig(
             path=self.path,
             seq_length=self.seq_length,

src/megatron/bridge/diffusion/models/flux/__init__.py

@@ -53,7 +53,7 @@
     get_flux_single_transformer_engine_spec,
 )
 from megatron.bridge.diffusion.models.flux.flux_model import Flux
-from megatron.bridge.diffusion.models.flux.flux_provider import FluxModelProvider14B, FluxProvider
+from megatron.bridge.diffusion.models.flux.flux_provider import FluxModelProvider12B, FluxProvider
 from megatron.bridge.diffusion.models.flux.layers import (
     EmbedND,
     MLPEmbedder,
@@ -66,7 +66,7 @@
 __all__ = [
     # Main model
     "Flux",
-    "FluxModelProvider14B",
+    "FluxModelProvider12B",
     "FluxProvider",
     # Transformer layers
     "MMDiTLayer",

src/megatron/bridge/diffusion/models/flux/flux_provider.py

@@ -135,8 +135,8 @@ def provide(self, pre_process=None, post_process=None, vp_stage=None) -> Flux:
 
 
 @dataclass
-class FluxModelProvider14B(FluxProvider):
-    """FLUX 14B model configuration.
+class FluxModelProvider12B(FluxProvider):
+    """FLUX 12B model configuration.
 
     Architecture: 19 joint + 38 single transformer blocks, hidden_size=3072,
     24 attention heads, context_dim=4096. Default seq_length=1024.

src/megatron/bridge/diffusion/recipes/flux/__init__.py

@@ -13,9 +13,9 @@
 # limitations under the License.
 
 from megatron.bridge.diffusion.recipes.flux.flux import (
-    flux_14b_finetune_config,
-    flux_14b_pretrain_config,
+    flux_12b_finetune_config,
+    flux_12b_pretrain_config,
 )
 
 
-__all__ = ["flux_14b_finetune_config", "flux_14b_pretrain_config"]
+__all__ = ["flux_12b_finetune_config", "flux_12b_pretrain_config"]

src/megatron/bridge/diffusion/recipes/flux/flux.py

@@ -18,7 +18,7 @@
 from megatron.core.distributed import DistributedDataParallelConfig
 
 from megatron.bridge.diffusion.data.flux.flux_energon_datamodule import FluxDatasetConfig
-from megatron.bridge.diffusion.models.flux.flux_provider import FluxModelProvider14B
+from megatron.bridge.diffusion.models.flux.flux_provider import FluxModelProvider12B
 from megatron.bridge.training.config import (
     CheckpointConfig,
     ConfigContainer,
@@ -30,9 +30,9 @@
 from megatron.bridge.training.mixed_precision import get_mixed_precision_config
 
 
-def flux_14b_pretrain_config() -> ConfigContainer:
+def flux_12b_pretrain_config() -> ConfigContainer:
     """
-    Return a pre-training configuration for FLUX model.
+    Return a pre-training configuration for FLUX 12B model.
 
     Default parallelism: TP=2, PP=1. Uses mock/synthetic data when data_paths is None.
     To customize (e.g. data paths, checkpoint dir), edit this recipe or add a new recipe
@@ -49,8 +49,8 @@ def flux_14b_pretrain_config() -> ConfigContainer:
     checkpoint_dir = os.path.join(run_output_dir, "checkpoints")
     tensorboard_dir = os.path.join(run_output_dir, "tb_logs")
 
-    # Model configuration - use FluxModelProvider14B (19 joint + 38 single layers)
-    model_cfg = FluxModelProvider14B()
+    # Model configuration - use FluxModelProvider12B (19 joint + 38 single layers)
+    model_cfg = FluxModelProvider12B()
     model_cfg.tensor_model_parallel_size = 2
     model_cfg.pipeline_model_parallel_size = 1
     model_cfg.pipeline_dtype = torch.bfloat16
@@ -145,14 +145,14 @@ def flux_14b_pretrain_config() -> ConfigContainer:
     return cfg
 
 
-def flux_14b_finetune_config(pretrained_checkpoint: str | None = None) -> ConfigContainer:
+def flux_12b_finetune_config(pretrained_checkpoint: str | None = None) -> ConfigContainer:
     """
-    Return a fine-tuning configuration for FLUX model.
+    Return a fine-tuning configuration for FLUX 12B model.
 
-    Uses the same defaults as flux_14b_pretrain_config() and overrides checkpoint to load from
+    Uses the same defaults as flux_12b_pretrain_config() and overrides checkpoint to load from
     pretrained_checkpoint when provided.
     """
-    cfg = flux_14b_pretrain_config()
+    cfg = flux_12b_pretrain_config()
     base_output_dir = os.path.join(os.getcwd(), "nemo_experiments")
     run_output_dir = os.path.join(base_output_dir, "default")
     checkpoint_dir = os.path.join(run_output_dir, "checkpoints")

tests/functional_tests/diffusion/flux/test_flux_pretrain.py

@@ -14,7 +14,7 @@
 
 """Functional smoke tests for Mcore FLUX pretrain mock runs.
 
-Uses the generic run_recipe.py entry point with flux_14b_pretrain_config and flux_step.
+Uses the generic run_recipe.py entry point with flux_12b_pretrain_config and flux_step.
 Mock/synthetic data is used when dataset.path is not set (no --mock flag).
 """
 
@@ -57,7 +57,7 @@ def test_flux_pretrain_mock(self, tmp_path):
             "--parallel-mode",
             "scripts/training/run_recipe.py",
             "--recipe",
-            "flux_14b_pretrain_config",
+            "flux_12b_pretrain_config",
             "--step_func",
             "flux_step",
             "model.tensor_model_parallel_size=1",

tests/unit_tests/diffusion/data/flux/test_flux_energon_datamodule.py

@@ -24,35 +24,6 @@
 pytestmark = [pytest.mark.unit]
 
 
-# -----------------------------------------------------------------------------
-# _is_empty_path
-# -----------------------------------------------------------------------------
-
-
-class TestIsEmptyPath:
-    """Tests for _is_empty_path helper."""
-
-    def test_none_is_empty(self):
-        assert flux_dm_mod._is_empty_path(None) is True
-
-    def test_empty_string_is_empty(self):
-        assert flux_dm_mod._is_empty_path("") is True
-
-    def test_whitespace_only_string_is_empty(self):
-        assert flux_dm_mod._is_empty_path("   ") is True
-
-    def test_non_empty_string_is_not_empty(self):
-        assert flux_dm_mod._is_empty_path("/path/to/data") is False
-        assert flux_dm_mod._is_empty_path("x") is False
-
-    def test_empty_list_is_empty(self):
-        assert flux_dm_mod._is_empty_path([]) is True
-
-    def test_non_empty_list_is_not_empty(self):
-        assert flux_dm_mod._is_empty_path(["/a"]) is False
-        assert flux_dm_mod._is_empty_path(["/a", "/b"]) is False
-
-
 # -----------------------------------------------------------------------------
 # FluxDatasetConfig (unified config: mock vs real at runtime)
 # -----------------------------------------------------------------------------
@@ -98,6 +69,9 @@ def test_build_datasets_empty_string_path_uses_mock_branch(self):
         fake_train, fake_val, fake_test = iter(["a"]), iter(["b"]), iter(["c"])
 
         class FakeMockConfig:
+            def __init__(self, **kwargs):
+                pass
+
             def build_datasets(self, context):
                 return (fake_train, fake_val, fake_test)
 
@@ -130,28 +104,6 @@ def test_build_datasets_string_path_uses_real_branch(self, monkeypatch):
         # Verify real config was built with the string path (via _FakeDiffusionDataModule)
         assert cfg.path == "/path/to/wds"
 
-    def test_build_datasets_list_path_uses_first_element(self, monkeypatch):
-        """When path is a list, build_datasets uses path[0] as path_str for FluxDataModuleConfig."""
-        received_paths = []
-
-        class CapturePathFake(_FakeDiffusionDataModule):
-            def __init__(self, *, path, **kwargs):
-                received_paths.append(path)
-                super().__init__(path=path, **kwargs)
-
-        monkeypatch.setattr(flux_dm_mod, "DiffusionDataModule", CapturePathFake)
-        cfg = flux_dm_mod.FluxDatasetConfig(
-            path=["/first/shard", "/second/shard"],
-            seq_length=1024,
-            packing_buffer_size=4,
-            micro_batch_size=2,
-            global_batch_size=8,
-            num_workers=0,
-        )
-        train, val, test = cfg.build_datasets(context=None)
-        assert list(train) == list("train")
-        assert received_paths == ["/first/shard"]
-
 
 # -----------------------------------------------------------------------------
 # FluxDataModuleConfig (real data only)

tests/unit_tests/diffusion/recipes/flux/test_flux_recipe.py

@@ -16,7 +16,7 @@
 
 from megatron.bridge.diffusion.data.flux.flux_energon_datamodule import FluxDatasetConfig
 from megatron.bridge.diffusion.models.flux.flux_provider import FluxProvider
-from megatron.bridge.diffusion.recipes.flux.flux import flux_14b_pretrain_config
+from megatron.bridge.diffusion.recipes.flux.flux import flux_12b_pretrain_config
 from megatron.bridge.training.config import ConfigContainer
 
 
@@ -28,7 +28,7 @@ class TestPretrainConfig:
 
     def test_pretrain_config_returns_complete_config(self):
         """Test that pretrain_config returns a ConfigContainer with all required components."""
-        config = flux_14b_pretrain_config()
+        config = flux_12b_pretrain_config()
 
         assert isinstance(config, ConfigContainer)
         assert isinstance(config.model, FluxProvider)
@@ -44,23 +44,23 @@ def test_pretrain_config_returns_complete_config(self):
 
     def test_pretrain_config_directory_structure(self):
         """Test that pretrain_config uses default directory structure."""
-        config = flux_14b_pretrain_config()
+        config = flux_12b_pretrain_config()
 
         assert "default" in config.checkpoint.save
         assert "default" in config.logger.tensorboard_dir
         assert config.checkpoint.save.endswith("checkpoints")
 
     def test_pretrain_config_default_training_parameters(self):
         """Test pretrain_config default training parameters."""
-        config = flux_14b_pretrain_config()
+        config = flux_12b_pretrain_config()
 
         assert config.train.train_iters == 10000
         assert config.train.global_batch_size == 16
         assert config.train.micro_batch_size == 1
 
     def test_pretrain_config_default_model_parameters(self):
         """Test that default model parameters are set correctly."""
-        config = flux_14b_pretrain_config()
+        config = flux_12b_pretrain_config()
 
         assert config.model.num_joint_layers == 19
         assert config.model.hidden_size == 3072
@@ -69,17 +69,17 @@ def test_pretrain_config_default_model_parameters(self):
 
     def test_pretrain_config_default_dataset_configuration(self):
         """Test pretrain_config default dataset parameters."""
-        config = flux_14b_pretrain_config()
+        config = flux_12b_pretrain_config()
 
         assert config.dataset.image_H == 1024
         assert config.dataset.image_W == 1024
         assert config.dataset.latent_channels == 16
 
     def test_pretrain_config_dataset_accepts_path_list(self):
         """Test that dataset config can be overridden to use real data paths."""
-        config = flux_14b_pretrain_config()
+        config = flux_12b_pretrain_config()
         assert config.dataset.path is None
 
-        # FluxDatasetConfig accepts path as list; recipe default is None
-        config.dataset.path = ["/some/data/path"]
-        assert config.dataset.path == ["/some/data/path"]
+        # FluxDatasetConfig accepts path as str; recipe default is None
+        config.dataset.path = "/some/data/path"
+        assert config.dataset.path == "/some/data/path"