[Flux] Fix missing field in flux argparse (#1120)

## Context
1. Bug fix. Fix missing fields in `arg_parser`.
2. Rebase to main folder changes,  add `ft_pg` field in train.py

Author: wwwjn

torchtitan/experiments/flux/dataset/flux_dataset.py

@@ -262,9 +262,13 @@ def build_flux_dataloader(
     ds = FluxDataset(
         dataset_name=dataset_name,
         dataset_path=dataset_path,
-        t5_tokenizer=FluxTokenizer(t5_encoder_name, max_length=max_t5_encoding_len),
+        t5_tokenizer=FluxTokenizer(
+            t5_encoder_name,
+            max_length=max_t5_encoding_len,
+        ),
         clip_tokenizer=FluxTokenizer(
-            clip_encoder_name, max_length=77
+            clip_encoder_name,
+            max_length=77,
         ),  # fix max_length for CLIP
         job_config=job_config,
         dp_rank=dp_rank,

torchtitan/experiments/flux/dataset/tokenizer.py

@@ -23,20 +23,20 @@ class FluxTokenizer(Tokenizer):
 
     """
 
-    def __init__(self, model_path: str = "t5-small", max_length: int = 77):
+    def __init__(self, model_path: str = "t5-small", max_length: int = 77, **hf_kwargs):
         super().__init__()
         self._n_words = 8  # TODO(jianiw): check
         self._max_length = max_length
 
-        self.is_clip = model_path.startswith("openai")
+        self.is_clip = "clip" in model_path.lower()
 
         if self.is_clip:
             self._tokenizer: CLIPTokenizer = CLIPTokenizer.from_pretrained(
-                model_path, max_length=max_length
+                model_path, max_length=max_length, **hf_kwargs
             )
         else:
             self._tokenizer: T5Tokenizer = T5Tokenizer.from_pretrained(
-                model_path, max_length=max_length
+                model_path, max_length=max_length, **hf_kwargs
             )
 
     def encode(

torchtitan/experiments/flux/flux_argparser.py

@@ -6,8 +6,6 @@
 
 import argparse
 
-import torch
-
 
 def extend_parser(parser: argparse.ArgumentParser) -> None:
     parser.add_argument(
@@ -26,31 +24,28 @@ def extend_parser(parser: argparse.ArgumentParser) -> None:
         "--encoder.t5_encoder",
         type=str,
         default="google/t5-v1_1-small",
-        help="T5 encoder to use, HuggingFace model name.",
+        help="T5 encoder to use, HuggingFace model name. This field could be either a local folder path, \
+        or a Huggingface repo name.",
     )
     parser.add_argument(
         "--encoder.clip_encoder",
         type=str,
         default="openai/clip-vit-large-patch14",
-        help="Clip encoder to use, HuggingFace model name.",
+        help="Clip encoder to use, HuggingFace model name. This field could be either a local folder path, \
+        or a Huggingface repo name.",
     )
     parser.add_argument(
-        "--encoder.encoder_dtype",
-        type=torch.dtype,
-        default=torch.bfloat16,
-        help="Which dtype to load for autoencoder. ",
+        "--encoder.autoencoder_path",
+        type=str,
+        default="torchtitan/experiments/flux/assets/autoencoder/ae.safetensors",
+        help="Autoencoder checkpoint path to load. This should be a local path referring to a safetensors file.",
     )
     parser.add_argument(
         "--encoder.max_t5_encoding_len",
         type=int,
         default=512,
         help="Maximum length of the T5 encoding.",
     )
-    parser.add_argument(
-        "--encoder.offload_encoder",
-        action="store_true",
-        help="Whether to shard the encoder using FSDP",
-    )
     # eval configs
     parser.add_argument(
         "--eval.enable_classifer_free_guidance",

torchtitan/experiments/flux/model/hf_embedder.py

@@ -11,7 +11,7 @@
 class FluxEmbedder(nn.Module):
     def __init__(self, version: str, **hf_kwargs):
         super().__init__()
-        self.is_clip = version.startswith("openai")
+        self.is_clip = "clip" in version.lower()
         self.output_key = "pooler_output" if self.is_clip else "last_hidden_state"
 
         if self.is_clip:

torchtitan/experiments/flux/tests/test_generate_image.py

@@ -79,26 +79,27 @@ def test_generate_image(self):
             device=torch_device,
             dtype=torch.bfloat16,
         )
-        clip_tokenizer = FluxTokenizer(
-            model_path=config.encoder.clip_encoder, max_length=77
-        )
+        t5_encoder = FluxEmbedder(
+            version=config.encoder.t5_encoder,
+        ).to(torch_device, dtype=torch.bfloat16)
         t5_tokenizer = FluxTokenizer(
             model_path=config.encoder.t5_encoder,
             max_length=config.encoder.max_t5_encoding_len,
         )
-        clip_encoder = FluxEmbedder(version=config.encoder.clip_encoder).to(
-            torch_device, dtype=torch.bfloat16
-        )
-        t5_encoder = FluxEmbedder(version=config.encoder.t5_encoder).to(
-            torch_device, dtype=torch.bfloat16
+        clip_encoder = FluxEmbedder(
+            version=config.encoder.clip_encoder,
+        ).to(torch_device, dtype=torch.bfloat16)
+        clip_tokenizer = FluxTokenizer(
+            model_path=config.encoder.clip_encoder,
+            max_length=77,
         )
 
         if torch.cuda.is_available():
             torch.cuda.synchronize()
         t1 = time.perf_counter()
 
         model = self._get_test_model(
-            context_in_dim=768, device=torch_device, dtype=torch.bfloat16
+            context_in_dim=4096, device=torch_device, dtype=torch.bfloat16
         )
         model.eval()
 

torchtitan/experiments/flux/train.py

@@ -51,17 +51,18 @@ def __init__(self, job_config: JobConfig):
         model_config = self.train_spec.config[job_config.model.flavor]
 
         self.autoencoder = load_ae(
-            job_config.encoder.auto_encoder_path,
+            job_config.encoder.autoencoder_path,
             model_config.autoencoder_params,
             device=self.device,
             dtype=self._dtype,
         )
-        self.clip_encoder = FluxEmbedder(version=job_config.encoder.clip_encoder).to(
-            device=self.device, dtype=self._dtype
-        )
-        self.t5_encoder = FluxEmbedder(version=job_config.encoder.t5_encoder).to(
-            device=self.device, dtype=self._dtype
-        )
+
+        self.clip_encoder = FluxEmbedder(
+            version=job_config.encoder.clip_encoder,
+        ).to(device=self.device, dtype=self._dtype)
+        self.t5_encoder = FluxEmbedder(
+            version=job_config.encoder.t5_encoder,
+        ).to(device=self.device, dtype=self._dtype)
 
         # Apply FSDP to the T5 model / CLIP model
         self.t5_encoder, self.clip_encoder = parallelize_encoders(
@@ -159,9 +160,10 @@ def train_step(self, input_dict: dict[str, torch.Tensor], labels: torch.Tensor):
             or parallel_dims.cp_enabled
         ):
             loss = loss.detach()
+            ft_pg = self.ft_manager.replicate_pg if self.ft_manager.enabled else None
             global_avg_loss, global_max_loss = (
-                dist_utils.dist_mean(loss, world_mesh["dp_cp"]),
-                dist_utils.dist_max(loss, world_mesh["dp_cp"]),
+                dist_utils.dist_mean(loss, world_mesh["dp_cp"], ft_pg),
+                dist_utils.dist_max(loss, world_mesh["dp_cp"], ft_pg),
             )
         else:
             global_avg_loss = global_max_loss = loss.item()

torchtitan/experiments/flux/train_configs/debug_model.toml

@@ -46,7 +46,7 @@ img_size = 256
 t5_encoder = "google/t5-v1_1-xxl"
 clip_encoder = "openai/clip-vit-large-patch14"
 max_t5_encoding_len = 4096
-auto_encoder_path = "torchtitan/experiments/flux/assets/autoencoder/ae.safetensors"  # Autoencoder to use for image
+autoencoder_path = "torchtitan/experiments/flux/assets/autoencoder/ae.safetensors"  # Autoencoder to use for image
 
 [eval]
 enable_classifer_free_guidance = true

torchtitan/experiments/flux/train_configs/flux_dev_model.toml

@@ -45,7 +45,7 @@ img_size = 256
 t5_encoder = "google/t5-v1_1-xxl"
 clip_encoder = "openai/clip-vit-large-patch14"
 max_t5_encoding_len = 4096
-auto_encoder_path = "torchtitan/experiments/flux/assets/autoencoder/ae.safetensors"  # Autoencoder to use for image
+autoencoder_path = "torchtitan/experiments/flux/assets/autoencoder/ae.safetensors"  # Autoencoder to use for image
 
 [eval]
 enable_classifer_free_guidance = true

torchtitan/experiments/flux/train_configs/flux_schnell_model.toml

@@ -45,7 +45,7 @@ img_size = 256
 t5_encoder = "google/t5-v1_1-xxl"
 clip_encoder = "openai/clip-vit-large-patch14"
 max_t5_encoding_len = 4096
-auto_encoder_path = "torchtitan/experiments/flux/assets/autoencoder/ae.safetensors"  # Autoencoder to use for image
+autoencoder_path = "torchtitan/experiments/flux/assets/autoencoder/ae.safetensors"  # Autoencoder to use for image
 
 [eval]
 enable_classifer_free_guidance = true