[Flux] Fix parallel dimension names (#1199)

## Context
- Currently we did not enable CP for flux model training, so we need to
remove the
- For FSDP/ HSDP, the correct dimension name should be ("dp_replicate",
"dp_shard") or ("dp_shard",)

Author: wwwjn

torchtitan/experiments/flux/parallelize_flux.py

@@ -28,13 +28,11 @@ def parallelize_flux(
     if job_config.activation_checkpoint.mode != "none":
         apply_ac(model, job_config.activation_checkpoint)
 
-    if (
-        parallel_dims.dp_shard_enabled or parallel_dims.dp_replicate_enabled
-    ):  # apply FSDP or HSDP
+    if parallel_dims.dp_shard_enabled:  # apply FSDP or HSDP
         if parallel_dims.dp_replicate_enabled:
-            dp_mesh_dim_names = ("dp_replicate", "dp_shard_cp")
+            dp_mesh_dim_names = ("dp_replicate", "dp_shard")
         else:
-            dp_mesh_dim_names = ("dp_shard_cp",)
+            dp_mesh_dim_names = ("dp_shard",)
 
         apply_fsdp(
             model,
@@ -122,13 +120,11 @@ def parallelize_encoders(
     parallel_dims: ParallelDims,
     job_config: JobConfig,
 ):
-    if (
-        parallel_dims.dp_shard_enabled or parallel_dims.dp_replicate_enabled
-    ):  # apply FSDP or HSDP
+    if parallel_dims.dp_shard_enabled:  # apply FSDP or HSDP
         if parallel_dims.dp_replicate_enabled:
-            dp_mesh_dim_names = ("dp_replicate", "dp_shard_cp")
+            dp_mesh_dim_names = ("dp_replicate", "dp_shard")
         else:
-            dp_mesh_dim_names = ("dp_shard_cp",)
+            dp_mesh_dim_names = ("dp_shard",)
 
         mp_policy = MixedPrecisionPolicy(
             param_dtype=TORCH_DTYPE_MAP[job_config.training.mixed_precision_param],