model fragments for diloco

Test Plan:
```bash
$ LOGLEVEL=INFO NCCL_DEBUG_SUBSYS=ALL NCCL_DEBUG=WARN TORCH_CPP_LOG_LEVEL=INFO CUDA_LAUNCH_BLOCKING=1 CUDA_VISIBLE_DEVICES=0,1 NGPU=2 ./run_train.sh --fault_tolerance.enable --fault_tolerance.group_size=1 --fault_tolerance.replica_id=0 --training.local_batch_size=2 --fault_tolerance.sync_steps=10 --fault_tolerance.semi_sync_method=diloco --parallelism.data_parallel_shard_degree=2 --fault_tolerance.num_fragments=2
```

Author: tushar00jain

torchtitan/config_manager.py

@@ -691,6 +691,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_module_names_per_stage",
     "pipeline_module_split",
+    "module_split",
 ]
 
 
@@ -333,3 +334,95 @@ 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)
+                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 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,10 +10,11 @@
 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.parallelize import parallelize_llama
 from .infra.pipeline import pipeline_llama
+from .infra.fault_tolerance import fragment_llama
 from .model.args import TransformerModelArgs
 from .model.model import Transformer
 from .model.state_dict_adapter import Llama3StateDictAdapter
@@ -70,12 +71,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,47 @@
+# 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_manager import JobConfig
+from torchtitan.distributed.pipeline import (
+    generate_module_names_per_stage,
+    module_split,
+)
+from torchtitan.tools.logging import logger
+
+from ..model.args import TransformerModelArgs
+
+def fragment_llama(
+    model: nn.Module,
+    job_config: JobConfig,
+    model_config: TransformerModelArgs,
+) -> list[nn.Module]:
+    ft = job_config.fault_tolerance
+
+    assert ft.num_fragments > 0
+
+    module_names_per_stage = ft.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.num_fragments == 1:
+            return [model]
+
+        module_names_per_stage = generate_module_names_per_stage(
+            ft.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

@@ -104,6 +104,16 @@ 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

@@ -8,7 +8,7 @@
 import os
 import time
 from datetime import timedelta
-from typing import Any, Generator, Iterable, Optional
+from typing import Any, Generator, Iterable, Optional, cast
 
 import torch
 from torch.distributed.elastic.multiprocessing.errors import record
@@ -42,6 +42,7 @@ class Trainer(torch.distributed.checkpoint.stateful.Stateful):
     # swappable training components in TrainSpec
     dataloader: train_spec_module.BaseDataLoader
     model_parts: list[torch.nn.Module]
+    ft_model_parts: list[torch.nn.Module]
     loss_fn: train_spec_module.LossFunction
     optimizers: train_spec_module.OptimizersContainer
     lr_schedulers: train_spec_module.LRSchedulersContainer
@@ -256,6 +257,16 @@ def __init__(self, job_config: JobConfig):
 
             self.model_parts = [model]
 
+            ft = job_config.fault_tolerance
+
+            if ft.enable:
+                train_spec = cast(train_spec_module.FaultTolerantTrainSpec, self.train_spec)
+                if train_spec.fragment_fn:
+                    self.ft_model_parts = train_spec.fragment_fn(model, job_config, model_args)
+                else:
+                    self.ft_model_parts = [model]
+
+
         self.ft_manager.maybe_set_all_reduce_hook(self.model_parts)
 
         # initialize device memory monitor and get peak flops for MFU calculation
@@ -501,7 +512,7 @@ def train(self):
             maybe_semi_sync_training(
                 job_config.fault_tolerance,
                 ft_manager=self.ft_manager,
-                model_parts=self.model_parts,
+                model_parts=self.ft_model_parts,
                 optimizer=self.optimizers,
             ),
         ):