refactor ParallelDims and CheckpointManager

Author: tianyu-l

scripts/estimate/estimation.py

@@ -48,7 +48,6 @@ def estimate_memory(job_config: JobConfig):
         pp=parallelism_config.pipeline_parallel_degree,
         ep=parallelism_config.expert_parallel_degree,
         world_size=world_size,
-        enable_loss_parallel=not parallelism_config.disable_loss_parallel,
     )
 
     # only FSDP and HSDP are supported
@@ -76,14 +75,14 @@ def estimate_memory(job_config: JobConfig):
 
     train_spec = get_train_spec(job_config.model.name)
 
-    # build meshes
-    world_mesh = parallel_dims.build_mesh(device_type="cuda")
-
     # build tokenizer
     tokenizer = train_spec.build_tokenizer_fn(job_config)
 
+    loss_parallel_enabled = (
+        parallel_dims.tp_enabled and not parallelism_config.disable_loss_parallel
+    )
     train_context = dist_utils.get_train_context(
-        parallel_dims.loss_parallel_enabled,
+        loss_parallel_enabled,
         job_config.parallelism.enable_compiled_autograd,
     )
 
@@ -108,7 +107,7 @@ def estimate_memory(job_config: JobConfig):
         model_converters.convert(model)
 
         # apply PT-D DP/TP parallelisms and activation checkpointing
-        train_spec.parallelize_fn(model, world_mesh, parallel_dims, job_config)
+        train_spec.parallelize_fn(model, parallel_dims, job_config)
 
         model.to_empty(device="cuda")
         if not active_fake_mode():
@@ -117,9 +116,7 @@ def estimate_memory(job_config: JobConfig):
 
         # build optimizer after applying parallelisms to the model
         ft_manager = init_ft_manager(job_config)
-        optimizers = build_optimizers(
-            [model], job_config, parallel_dims, world_mesh, ft_manager
-        )
+        optimizers = build_optimizers([model], job_config, parallel_dims, ft_manager)
         lr_schedulers = build_lr_schedulers(optimizers.optimizers, job_config)
         # Post optimizer step model converters hook.
         # e.g. calculate float8 dynamic amax/scale for all-parameter for FSDP2

scripts/generate/test_generate.py

@@ -127,14 +127,12 @@ def test_generate(
             pp=1,
             ep=1,
             world_size=world_size,
-            enable_loss_parallel=False,
         )
-        # Build world mesh for parallelism
-        world_mesh = parallel_dims.build_mesh(device_type=device_type)
+        world_mesh = parallel_dims.world_mesh
 
         # apply_tp (with Sequence Parallel) on unevenly sharded
         # sequences would require https://github.com/pytorch/torchtitan/pull/686
-        apply_tp_minus_sp(model, world_mesh["tp"])
+        apply_tp_minus_sp(model, parallel_dims.world_mesh["tp"])
 
     dist_utils.set_determinism(world_mesh, device, seed, deterministic)
 

tests/unit_tests/test_model_converter.py

@@ -23,7 +23,6 @@ def build_parallel_dims(job_config, world_size):
         pp=parallelism_config.pipeline_parallel_degree,
         ep=parallelism_config.expert_parallel_degree,
         world_size=world_size,
-        enable_loss_parallel=not parallelism_config.disable_loss_parallel,
     )
     return parallel_dims
 

torchtitan/components/checkpoint.py

@@ -26,8 +26,8 @@
 )
 from torch.distributed.checkpoint.state_dict_saver import AsyncCheckpointerType
 from torch.distributed.checkpoint.stateful import Stateful
-from torch.utils.data import DataLoader
 
+from torchtitan.components.dataloader import BaseDataLoader
 from torchtitan.components.ft import FTManager
 from torchtitan.components.lr_scheduler import LRSchedulersContainer
 from torchtitan.components.optimizer import OptimizersContainer
@@ -180,17 +180,19 @@ class CheckpointManager:
 
     def __init__(
         self,
-        dataloader: DataLoader,
+        dataloader: BaseDataLoader | None,
         model_parts: list[nn.Module],
         optimizers: OptimizersContainer,
         lr_schedulers: LRSchedulersContainer,
         states: dict[str, Any],
         job_config: JobConfig,
-        ft_manager: FTManager,
+        ft_manager: FTManager | None = None,
     ) -> None:
         ckpt_config = job_config.checkpoint
         self.enable_checkpoint = ckpt_config.enable_checkpoint
-        self.ft_manager = ft_manager.manager if ft_manager.enabled else None
+        self.ft_manager = (
+            ft_manager.manager if ft_manager and ft_manager.enabled else None
+        )
 
         if self.ft_manager:
             optimizers.init_cache_state_dict()

torchtitan/components/optimizer.py

@@ -15,7 +15,6 @@
     StateDictOptions,
 )
 from torch.distributed.checkpoint.stateful import Stateful
-from torch.distributed.device_mesh import DeviceMesh
 from torch.optim import Optimizer
 
 from torchtitan.components.ft import FTManager, has_torchft
@@ -244,7 +243,6 @@ def build_optimizers(
     model_parts: list[nn.Module],
     job_config: JobConfig,
     parallel_dims: ParallelDims,
-    world_mesh: DeviceMesh,
     ft_manager: FTManager,
 ) -> OptimizersContainer:
     """Create a OptimizersContainer for the given model parts and job config.

torchtitan/components/validate.py

@@ -50,14 +50,12 @@ def __init__(
         dp_rank: int,
         tokenizer: BaseTokenizer,
         parallel_dims: ParallelDims,
-        world_mesh: torch.distributed.DeviceMesh,
         loss_fn: LossFunction,
         validation_context: Generator[None, None, None],
         maybe_enable_amp: Generator[None, None, None],
     ):
         self.job_config = job_config
         self.parallel_dims = parallel_dims
-        self.world_mesh = world_mesh
         self.loss_fn = loss_fn
         self.validation_dataloader = build_hf_validation_dataloader(
             job_config=job_config,
@@ -78,6 +76,8 @@ def validate(
         model = model_parts[0]
         model.eval()
 
+        parallel_dims = self.parallel_dims
+
         accumulated_losses = []
         device_type = utils.device_type
         num_steps = 0
@@ -96,13 +96,13 @@ def validate(
 
             optional_context_parallel_ctx = (
                 dist_utils.create_context_parallel_ctx(
-                    cp_mesh=self.world_mesh["cp"],
+                    cp_mesh=parallel_dims.world_mesh["cp"],
                     cp_buffers=[inputs, labels] + [m.freqs_cis for m in model_parts],
                     cp_seq_dims=[1, 1] + [0 for _ in model_parts],
                     cp_no_restore_buffers={inputs, labels},
                     cp_rotate_method=self.job_config.parallelism.context_parallel_rotate_method,
                 )
-                if self.parallel_dims.cp_enabled
+                if parallel_dims.cp_enabled
                 else None
             )
 
@@ -119,8 +119,10 @@ def validate(
         # Compute average loss
         loss = torch.sum(torch.stack(accumulated_losses))
         loss /= num_steps
-        if self.parallel_dims.dp_cp_enabled:
-            global_avg_loss = dist_utils.dist_mean(loss, self.world_mesh["dp_cp"])
+        if parallel_dims.dp_cp_enabled:
+            global_avg_loss = dist_utils.dist_mean(
+                loss, parallel_dims.world_mesh["dp_cp"]
+            )
         else:
             global_avg_loss = loss
 
@@ -144,7 +146,6 @@ def build_validator(
     dp_rank: int,
     tokenizer: BaseTokenizer,
     parallel_dims: ParallelDims,
-    world_mesh: torch.distributed.DeviceMesh,
     loss_fn: LossFunction,
     validation_context: Generator[None, None, None],
     maybe_enable_amp: Generator[None, None, None],
@@ -156,7 +157,6 @@ def build_validator(
         dp_rank=dp_rank,
         tokenizer=tokenizer,
         parallel_dims=parallel_dims,
-        world_mesh=world_mesh,
         loss_fn=loss_fn,
         validation_context=validation_context,
         maybe_enable_amp=maybe_enable_amp,

torchtitan/distributed/parallel_dims.py

@@ -10,6 +10,7 @@
 from torch.distributed.device_mesh import DeviceMesh, init_device_mesh
 
 from torchtitan.tools.logging import logger
+from torchtitan.tools.utils import device_type
 
 
 __all__ = ["ParallelDims"]
@@ -24,7 +25,8 @@ class ParallelDims:
     pp: int
     ep: int
     world_size: int
-    enable_loss_parallel: bool
+
+    _world_mesh: DeviceMesh = None
 
     def __post_init__(self):
         self._validate()
@@ -55,16 +57,16 @@ def _validate(self):
             # EP would borrow all cp and some dp_shard degree
             assert ep % cp == 0 and (dp_shard * cp) % ep == 0
 
-    def build_mesh(self, device_type: str) -> DeviceMesh:
+    def build_mesh(self) -> DeviceMesh:
         # TODO: Current implementation of ParallelDims for dp2ep Expert Parallel
         #       is not very clean, due to the limited support from DeviceMesh
         #       for creating two staggered meshes. Will improve.
         if self.ep > 1:
-            return self._build_mesh_with_ep(device_type)
+            return self._build_mesh_with_ep()
         else:
-            return self._build_mesh_without_ep(device_type)
+            return self._build_mesh_without_ep()
 
-    def _build_mesh_with_ep(self, device_type: str) -> DeviceMesh:
+    def _build_mesh_with_ep(self) -> DeviceMesh:
         # With ep, dp_shard and ep are derived submeshes:
         # dp_shard = dp_shard_mod_ep * dp_shard_in_ep
         # ep = dp_shard_in_ep * cp
@@ -128,7 +130,7 @@ def _build_mesh_with_ep(self, device_type: str) -> DeviceMesh:
 
         return mesh
 
-    def _build_mesh_without_ep(self, device_type: str) -> DeviceMesh:
+    def _build_mesh_without_ep(self) -> DeviceMesh:
         dims = []
         names = []
         for d, name in zip(
@@ -173,6 +175,14 @@ def _build_mesh_without_ep(self, device_type: str) -> DeviceMesh:
 
         return mesh
 
+    @property
+    def world_mesh(self) -> str:
+        # doing late init so ParallelDims can still be used as a lightweight
+        # dataclass without having to initialize the world mesh
+        if self._world_mesh is None:
+            self._world_mesh = self.build_mesh()
+        return self._world_mesh
+
     @property
     def dp_enabled(self):
         return self.dp_replicate > 1 or self.dp_shard > 1
@@ -206,18 +216,24 @@ def pp_enabled(self):
         return self.pp > 1
 
     @property
-    def loss_parallel_enabled(self):
-        return self.tp > 1 and self.enable_loss_parallel
+    def ep_enabled(self):
+        return self.ep > 1
 
     @cached_property
     def non_data_parallel_size(self):
         return self.cp * self.tp * self.pp
 
-    @property
-    def ep_enabled(self):
-        return self.ep > 1
+    @cached_property
+    def seq_len_divisor(self):
+        # Sequence Parallel requires that seq_len be divisible by TP degree.
+        # https://github.com/pytorch/torchtitan/pull/640#discussion_r1849481001
 
-    @property
+        # Context Parallel requires that seq_len be divisible by 2 * CP degree,
+        # when load balancing is enabled (by default).
+        # https://github.com/pytorch/pytorch/blob/4f62dcc/torch/distributed/tensor/experimental/_attention.py#L1246
+        return self.tp * (self.cp * 2)
+
+    @cached_property
     def dense_params_mesh_ndim(self):
-        # Note: EP params mesh ndim is 1 more due to the 'ep' mesh
+        # Note: In dp2ep EP, EP params mesh ndim is 1 more due to the 'ep' mesh
         return self.dp_replicate_enabled + self.fsdp_enabled + self.tp_enabled

torchtitan/experiments/deepseek_v3/train_ds_real.py

@@ -155,8 +155,8 @@ def run_full_model(
         pp=pp_size,
         cp=1,
         tp=1,
+        ep=1,
         world_size=world_mesh.size(),
-        enable_loss_parallel=False,
     )
 
     metrics_processor = build_metrics_processor(
@@ -180,7 +180,7 @@ def run_full_model(
     loss_fn = cross_entropy_loss  # torch.nn.functional.cross_entropy
 
     ft_manager = ft.init_ft_manager(config)
-    optimizer = build_optimizers([model], config, ft_manager)
+    optimizer = build_optimizers([model], config, proxy_parallel_dims, ft_manager)
 
     lr_scheduler = build_lr_schedulers(optimizer, config)
 

torchtitan/experiments/flux/infra/parallelize.py

@@ -21,7 +21,6 @@
 
 def parallelize_flux(
     model: nn.Module,
-    world_mesh: DeviceMesh,
     parallel_dims: ParallelDims,
     job_config: JobConfig,
 ):
@@ -36,7 +35,7 @@ def parallelize_flux(
 
         apply_fsdp(
             model,
-            world_mesh[tuple(dp_mesh_dim_names)],
+            parallel_dims.world_mesh[tuple(dp_mesh_dim_names)],
             param_dtype=TORCH_DTYPE_MAP[job_config.training.mixed_precision_param],
             reduce_dtype=TORCH_DTYPE_MAP[job_config.training.mixed_precision_reduce],
             cpu_offload=job_config.training.enable_cpu_offload,
@@ -117,7 +116,6 @@ def apply_ac(model: nn.Module, ac_config):
 def parallelize_encoders(
     t5_model: nn.Module,
     clip_model: nn.Module,
-    world_mesh: DeviceMesh,
     parallel_dims: ParallelDims,
     job_config: JobConfig,
 ):
@@ -132,7 +130,7 @@ def parallelize_encoders(
             reduce_dtype=TORCH_DTYPE_MAP[job_config.training.mixed_precision_reduce],
         )
         fsdp_config = {
-            "mesh": world_mesh[tuple(dp_mesh_dim_names)],
+            "mesh": parallel_dims.world_mesh[tuple(dp_mesh_dim_names)],
             "mp_policy": mp_policy,
         }
         if job_config.training.enable_cpu_offload:

torchtitan/experiments/flux/train.py

@@ -36,7 +36,7 @@ def __init__(self, job_config: JobConfig):
         # (mainly for debugging, expect perf loss).
         # For Flux model, we need distinct seed across FSDP ranks to ensure we randomly dropout prompts info in dataloader
         dist_utils.set_determinism(
-            self.world_mesh,
+            self.parallel_dims.world_mesh,
             self.device,
             job_config.training.seed,
             job_config.training.deterministic,
@@ -77,7 +77,6 @@ def __init__(self, job_config: JobConfig):
         self.t5_encoder, self.clip_encoder = parallelize_encoders(
             t5_model=self.t5_encoder,
             clip_model=self.clip_encoder,
-            world_mesh=self.world_mesh,
             parallel_dims=self.parallel_dims,
             job_config=job_config,
         )