model fragments for diloco

Summary:
- add a configuration option for users to provide how they want to partition the model
- if this is provided, the model needs to implement `FaultTolerantTrainingSpec` that defines the framentation function to split the model based on the configuration
- determine the model fragments in training script to pass to ft manager

Test Plan:
Running llama3 8b parameters with 2 fragments, 1 step delay, each fragment gets synced every 20 steps

<img width="944" height="545" alt="image" src="https://github.com/user-attachments/assets/6d16f486-7260-49d6-8ba3-3e98cd331e58" />

Author: tushar00jain

torchtitan/components/ft.py

@@ -7,13 +7,18 @@
 import importlib
 from contextlib import nullcontext
 from datetime import timedelta
-from typing import ContextManager, Optional, TYPE_CHECKING, Union
+from typing import cast, ContextManager, Optional, TYPE_CHECKING, Union
 
 import torch
 import torch.distributed as dist
 from torch.distributed._composable.fsdp.fully_shard import FSDPModule
 from torch.distributed.distributed_c10d import ReduceOp
 from torchtitan.config.job_config import FaultTolerance as FTConfig
+from torchtitan.protocols.train_spec import (
+    BaseModelArgs,
+    FaultTolerantTrainSpec,
+    TrainSpec,
+)
 
 if importlib.util.find_spec("torchft") is not None:
     import torchft as ft
@@ -108,8 +113,10 @@ def loss_sync_pg(
 def maybe_semi_sync_training(
     ft_config: FTConfig,
     ft_manager: FTManager,
-    model_parts: list[torch.nn.Module],
+    model: torch.nn.Module,
+    model_args: BaseModelArgs,
     optimizer: torch.optim.Optimizer,
+    train_spec: TrainSpec,
 ) -> ContextManager[Union["local_sgd.DiLoCo", "local_sgd.LocalSGD", None]]:
     """
     If TorchFT is enabled and the config is set, use semi_sync_method
@@ -122,6 +129,12 @@ def maybe_semi_sync_training(
             ft_manager._manager is not None
         ), "FTManager must be enabled to use semi-sync training."
         if semi_sync_method.lower() == "diloco":
+            train_spec = cast(FaultTolerantTrainSpec, train_spec)
+            if train_spec.fragment_fn:
+                model_parts = train_spec.fragment_fn(model, ft_config, model_args)
+            else:
+                model_parts = [model]
+
             # Create the outer optimizer based on the inner optimizer parameters.
             outer_optimizers = []
             for model in model_parts:
@@ -142,10 +155,10 @@ def maybe_semi_sync_training(
                 fragment_update_alpha=ft_config.fragment_update_alpha,
             )
         elif semi_sync_method.lower() == "local_sgd":
-            assert len(model_parts) == 1
+            assert len(model) == 1
             return local_sgd.LocalSGD(
                 manager=ft_manager._manager,
-                model=model_parts[0],
+                model=model,
                 optimizer=optimizer,
                 sync_every=ft_config.sync_steps,
             )

torchtitan/config/job_config.py

@@ -686,6 +686,18 @@ class FaultTolerance:
     This is only used when "semi_sync_method" is set.
     """
 
+    module_names_per_model_chunk: list[list[str]] = field(default_factory=list)
+    """
+    Specify a list of lists containing the FQNs (Fully Qualified Names) of modules for each model chunk.
+    Each inner list represents one model chunk and contains the module names that belong to that chunk.
+    e.g. [['tok_embeddings', 'layers.0'], ['layers.1', 'layers.2'], ['layers.3', 'layers.4']]
+    will create 3 chunks: the first containing tok_embeddings and layers.0,
+    the second containing layers.1 and layers.2, and the third containing layers.3 and layers.4.
+    This provides more explicit control over which modules belong to each chunk compared to split points.
+    """
+
+    num_fragments: int = 1
+
 
 @dataclass
 class Experimental:

torchtitan/distributed/pipeline.py

@@ -30,6 +30,7 @@
     "stage_ids_this_rank",
     "generate_llm_fqn_per_model_part",
     "pipeline_module_split",
+    "module_split",
 ]
 
 
@@ -118,7 +119,7 @@ def stage_ids_this_rank(
             stages_per_rank == 2
         ), f"v schedules assume 2 stages per rank, got {stages_per_rank}"
         stage_v_pairs = list(
-            zip(range(pp_size), range(num_stages - 1, pp_size - 1, -1))
+            zip(range(pp_size), range(num_stages - 1, pp_size - 1, -1), strict=True)
         )
         return stage_v_pairs[pp_rank]
 
@@ -352,3 +353,96 @@ def _build_stage_from_modules(
         models.append(model_chunk)
 
     return stages, models
+
+
+def module_split(
+    model: nn.Module,
+    module_names_per_stage: list[list[str]],
+) -> list[nn.Module]:
+    """
+    This API creates pipeline stages based on specified module names for each stage.
+    This method updates the model in place.
+
+    Args:
+        model: The complete model to be split
+        module_names_per_stage: List of lists, where each inner list contains the module names
+                               that should be included in that stage. Module names should be
+                               dot-separated paths. Examples:
+                               - "tok_embeddings" for token embeddings
+                               - "layers.0", "layers.1" for specific transformer layers
+                               - "norm" for the final normalization layer
+                               - "output" for the output projection layer
+
+    Returns:
+        List of model chunks
+
+    Example usage:
+        module_names_per_stage = [
+            ["tok_embeddings", "layers.0"],     # Stage 0: embeddings + first layer
+            ["layers.1", "layers.2"],           # Stage 1: middle layers
+            ["norm", "output"]                  # Stage 2: final norm + output
+        ]
+    """
+
+    def _build_stage_from_modules(stage_idx: int, module_names: list[str]) -> nn.Module:
+        stage_model = nn.Module()
+        # Create a set of modules to keep for faster lookup
+        modules_to_keep = set(module_names)
+        print(f"Stage {stage_idx}: Modules to keep: {modules_to_keep}")
+        for module_name, module_value in model.named_children():
+            # Handle layer-like structures (e.g., "layers.0", "layers.1")
+            if isinstance(module_value, (nn.ModuleDict, nn.ModuleList)):
+                layers_to_keep = {
+                    name.split(".", 1)[1]
+                    for name in modules_to_keep
+                    if name.startswith(f"{module_name}.")
+                }
+
+                if not layers_to_keep:
+                    continue
+
+                # Keep only specified layers
+                if isinstance(module_value, nn.ModuleDict):
+                    for layer_name in list(module_value.keys()):
+                        if layer_name in layers_to_keep:
+                            setattr(
+                                stage_model,
+                                f"{module_name}.{layer_name}",
+                                module_value[layer_name],
+                            )
+                else:
+                    indices_to_keep = {
+                        int(idx) for idx in layers_to_keep if idx.isdigit()
+                    }
+                    new_layers = nn.ModuleList(
+                        [
+                            layer
+                            for i, layer in enumerate(module_value)
+                            if i in indices_to_keep
+                        ]
+                    )
+                    setattr(stage_model, module_name, new_layers)
+
+                continue
+
+            # Handle simple module attributes (e.g., "linear", "norm")
+            if module_name not in modules_to_keep:
+                continue
+
+            setattr(stage_model, module_name, module_value)
+
+        return stage_model
+
+    num_stages = len(module_names_per_stage)
+    models = []
+
+    for stage_idx in range(num_stages):
+        module_names = module_names_per_stage[stage_idx]
+        model_chunk = _build_stage_from_modules(
+            stage_idx,
+            module_names,
+        )
+        logger.info(f"building stage_idx {stage_idx} " f"with modules {module_names}")
+        models.append(model_chunk)
+
+    return models

torchtitan/models/llama3/__init__.py

@@ -10,7 +10,8 @@
 from torchtitan.components.tokenizer import build_hf_tokenizer
 from torchtitan.components.validate import build_validator
 from torchtitan.datasets.hf_datasets import build_hf_dataloader
-from torchtitan.protocols.train_spec import register_train_spec, TrainSpec
+from torchtitan.protocols.train_spec import FaultTolerantTrainSpec, register_train_spec
+from .infra.fault_tolerance import fragment_llama
 
 from .infra.parallelize import parallelize_llama
 from .infra.pipeline import pipeline_llama
@@ -71,12 +72,13 @@
 
 
 register_train_spec(
-    TrainSpec(
+    FaultTolerantTrainSpec(
         name="llama3",
         model_cls=Transformer,
         model_args=llama3_configs,
         parallelize_fn=parallelize_llama,
         pipelining_fn=pipeline_llama,
+        fragment_fn=fragment_llama,
         build_optimizers_fn=build_optimizers,
         build_lr_schedulers_fn=build_lr_schedulers,
         build_dataloader_fn=build_hf_dataloader,

torchtitan/models/llama3/infra/fault_tolerance.py

@@ -0,0 +1,43 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# This file is used to setup the model for fault tolerance
+
+import torch.nn as nn
+
+from torchtitan.config.job_config import FaultTolerance as FTConfig
+from torchtitan.distributed.pipeline import (
+    generate_llm_fqn_per_model_part,
+    module_split,
+)
+
+from ..model.args import TransformerModelArgs
+
+
+def fragment_llama(
+    model: nn.Module,
+    ft_config: FTConfig,
+    model_config: TransformerModelArgs,
+) -> list[nn.Module]:
+    assert ft_config.num_fragments > 0
+
+    module_names_per_stage = ft_config.module_names_per_model_chunk
+
+    input_weight = 1  # Weight for tok_embeddings
+    output_weight = 1  # Weight for norm + output layers
+
+    if module_names_per_stage == []:
+        if ft_config.num_fragments == 1:
+            return [model]
+
+        module_names_per_stage = generate_llm_fqn_per_model_part(
+            ft_config.num_fragments, model_config.n_layers, input_weight, output_weight
+        )
+
+    model_fragments = module_split(model, module_names_per_stage)
+    print(f"Created {len(model_fragments)} model fragments")
+
+    return model_fragments

torchtitan/protocols/train_spec.py

@@ -56,6 +56,14 @@ class TrainSpec:
     state_dict_adapter: type[StateDictAdapter] | None = None
 
 
+FragmentFunction: TypeAlias = Callable[..., list[nn.Module]]
+
+
+@dataclass
+class FaultTolerantTrainSpec(TrainSpec):
+    fragment_fn: FragmentFunction | None = None
+
+
 _train_specs = {}
 
 

torchtitan/train.py

@@ -38,6 +38,7 @@ class Trainer(torch.distributed.checkpoint.stateful.Stateful):
     job_config: JobConfig
     parallel_dims: ParallelDims
     train_spec: train_spec_module.TrainSpec
+    model_args: train_spec_module.BaseModelArgs
 
     # swappable training components in TrainSpec
     tokenizer: train_spec_module.BaseTokenizer | None
@@ -146,6 +147,7 @@ def __init__(self, job_config: JobConfig):
         model_args = self.train_spec.model_args[job_config.model.flavor]
         # set the model args from training job configs
         model_args.update_from_config(job_config)
+        self.model_args = model_args
 
         logger.info(
             f"Building {self.train_spec.name} {job_config.model.flavor} with {model_args}"
@@ -529,8 +531,10 @@ def train(self):
             maybe_semi_sync_training(
                 job_config.fault_tolerance,
                 ft_manager=self.ft_manager,
-                model_parts=self.model_parts,
+                model=self.model_parts[0],
+                model_args=self.model_args,
                 optimizer=self.optimizers,
+                train_spec=self.train_spec,
             ),
         ):
             data_iterator = self.batch_generator(self.dataloader)