refactor ParallelDims and CheckpointManager (#1384)

This PR does the following:
1. move `world_mesh` into `ParallelDims`, as they have a close
relationship
2. move `enable_loss_parallel` out of `ParallelDims` constructor
3. add a convenient property `seq_len_divisor` to `ParallelDims`
4. set `dataloader` and `ft_manager` as optional in `CheckpointManager`
5. some minor improvements on typing and code organization

Author: tianyu-l

scripts/estimate/estimation.py

@@ -39,6 +39,12 @@ def estimate_memory(job_config: JobConfig):
         job_config.training.compile = False
         job_config.parallelism.enable_compiled_autograd = False
 
+    # init fake pg
+    store = FakeStore()
+    torch.distributed.init_process_group(
+        "fake", rank=int(os.environ["LOCAL_RANK"]), world_size=world_size, store=store
+    )
+
     parallelism_config = job_config.parallelism
     parallel_dims = ParallelDims(
         dp_shard=parallelism_config.data_parallel_shard_degree,
@@ -48,8 +54,9 @@ 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,
     )
+    # ParallelDims.build_mesh has to happen outside of the FakeTensorMode
+    _ = parallel_dims.world_mesh
 
     # only FSDP and HSDP are supported
     if (
@@ -68,28 +75,21 @@ def estimate_memory(job_config: JobConfig):
     device = torch.device(f"cuda:{int(os.environ['LOCAL_RANK'])}")
     torch.cuda.set_device(device)
 
-    # init fake pg
-    store = FakeStore()
-    torch.distributed.init_process_group(
-        "fake", rank=int(os.environ["LOCAL_RANK"]), world_size=world_size, store=store
-    )
-
     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,
     )
 
     # build model (using meta init)
-    model_cls = train_spec.cls
-    model_args = train_spec.config[job_config.model.flavor]
+    model_args = train_spec.model_args[job_config.model.flavor]
     model_args.update_from_config(job_config, tokenizer)
 
     with (
@@ -101,14 +101,14 @@ def estimate_memory(job_config: JobConfig):
             f"Building {train_spec.name} {job_config.model.flavor} with {model_args}"
         )
         with torch.device("meta"):
-            model = model_cls(model_args)
+            model = train_spec.model_cls(model_args)
 
         # Build the collection of model converters. No-op if `model.converters` empty
         model_converters = build_model_converters(job_config, parallel_dims)
         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 +117,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

@@ -106,14 +106,13 @@ def test_generate(
     # Tokenizer setup
     tokenizer = train_spec.build_tokenizer_fn(config)
 
-    model_cls = train_spec.cls
-    model_args = train_spec.config[config.model.flavor]
+    model_args = train_spec.model_args[config.model.flavor]
     model_args.update_from_config(config, tokenizer)
 
     init_device = "meta" if world_size > 1 else device
     with torch.device(init_device):
         logger.info(f"Init model on init_device: {init_device}")
-        model = model_cls(model_args)
+        model = train_spec.model_cls(model_args)
 
     world_mesh = None
     # Init distributed env
@@ -127,14 +126,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
 

tests/unit_tests/test_train_spec.py

@@ -9,12 +9,14 @@
 import pytest
 import torch
 import torch.nn as nn
+from torchtitan.components.ft import FTManager
 from torchtitan.components.loss import build_cross_entropy_loss
 from torchtitan.components.lr_scheduler import build_lr_schedulers
 from torchtitan.components.optimizer import build_optimizers, OptimizersContainer
 from torchtitan.components.tokenizer import build_hf_tokenizer
 from torchtitan.config_manager import JobConfig
 from torchtitan.datasets.hf_datasets import build_hf_dataloader
+from torchtitan.distributed.parallel_dims import ParallelDims
 from torchtitan.models.llama3 import parallelize_llama, pipeline_llama
 from torchtitan.protocols.train_spec import (
     apply_to_train_specs,
@@ -39,7 +41,10 @@ def init_weights(self, buffer_device: torch.device | None = None) -> None:
 
 
 def fake_build_optimizers(
-    model_parts: list[nn.Module], job_config: JobConfig
+    model_parts: list[nn.Module],
+    job_config: JobConfig,
+    parallel_dims: ParallelDims,
+    ft_manager: FTManager,
 ) -> OptimizersContainer:
     optimizer_kwargs = {
         "lr": 0.1,
@@ -57,11 +62,11 @@ def fake_build_optimizers(
 
 class TestTrainSpec:
     def test_register_train_spec(self):
-        fake_config = {"fake": None}
+        fake_config = {"fake": BaseModelArgs()}
         spec = TrainSpec(
             name="fake",
-            cls=FakeModel,
-            config=fake_config,
+            model_cls=FakeModel,
+            model_args=fake_config,
             parallelize_fn=parallelize_llama,
             pipelining_fn=pipeline_llama,
             build_optimizers_fn=build_optimizers,
@@ -78,11 +83,11 @@ def test_register_train_spec(self):
             new_spec = get_train_spec("fake2")
 
     def test_optim_hook(self):
-        fake_config = {"fake": None}
+        fake_config = {"fake": BaseModelArgs()}
         spec = TrainSpec(
             name="fake2",
-            cls=FakeModel,
-            config=fake_config,
+            model_cls=FakeModel,
+            model_args=fake_config,
             parallelize_fn=parallelize_llama,
             pipelining_fn=pipeline_llama,
             build_optimizers_fn=fake_build_optimizers,
@@ -111,21 +116,27 @@ def register_optimizer_hook_to_spec(spec: TrainSpec) -> TrainSpec:
             original_build_optimizers_fn = spec.build_optimizers_fn
 
             def my_build_optimizer_fn(
-                model_parts: list[nn.Module], job_config: JobConfig
+                model_parts: list[nn.Module],
+                job_config: JobConfig,
+                parallel_dims: ParallelDims,
+                ft_manager: FTManager,
             ) -> OptimizersContainer:
-                optimizers = original_build_optimizers_fn(model_parts, job_config)
+                optimizers = original_build_optimizers_fn(
+                    model_parts, job_config, parallel_dims, ft_manager
+                )
                 optimizers.register_step_post_hook(
                     partial(my_hook, model_parts=model_parts)
                 )
                 return optimizers
 
             spec.build_optimizers_fn = my_build_optimizer_fn
+            return spec
 
         apply_to_train_specs(register_optimizer_hook_to_spec)
 
-        model = new_spec.cls(BaseModelArgs())
+        model = new_spec.model_cls(BaseModelArgs())
         model_parts = [model]
-        optimizers = new_spec.build_optimizers_fn(model_parts, JobConfig())
+        optimizers = new_spec.build_optimizers_fn(model_parts, None, None, None)
         assert optimizers.optimizers[0].__class__.__name__ == "Adam"
         batch = torch.randn(8, 8)
         model(batch).sum().backward()

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/tokenizer.py

@@ -7,17 +7,15 @@
 
 import json
 
-import logging
 import os
 from abc import ABC, abstractmethod
 from typing import Any, Optional, Union
 
 from tokenizers import AddedToken, Tokenizer
 from torchtitan.config_manager import JobConfig
+from torchtitan.tools.logging import logger
 from typing_extensions import override
 
-logger = logging.getLogger(__name__)
-
 
 class BaseTokenizer(ABC):
     # base tokenizer interface, for typing purpose mainly

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,