Update

[ghstack-poisoned]

Author: fegin

.gitignore

@@ -5,6 +5,7 @@ __pycache__
 build
 outputs
 dist/*
+.vscode
 
 # data
 data

README.md

@@ -52,7 +52,7 @@ Our guiding principles when building `torchtitan`:
     - estimate FSDP/HSDP memory usage without materializing the model
     - run distributed inference with Tensor Parallel
 
-We report our [Performance](docs/performance.md) verified on 64/128 GPUs.
+We report [performance](docs/performance.md) on up to 512 GPUs, and verify [loss converging](docs/converging.md) correctness of various techniques.
 
 ### Dive into the code
 

assets/images/llama3_1_405B_loss_curves.png

@@ -35,27 +35,25 @@ If the technique is not a parallelism
 
 ## Example
 
-Setup
+This is a series of loss-converging tests covering both parallelisms and training optimizations.
+Results are obtained on 2025/01/21, with the latest `torch`, `torchao`, and `torchtitan`.
+
+### Setup
 - Base config: [train_configs/llama3_8b.toml](../train_configs/llama3_8b.toml)
-- `training.batch_size = 4`
-  - a minimum for Pipeline Parallel with `pipeline_parallel_degree = 2` and `pipeline_parallel_schedule = "Interleaved1F1B"`
+- `training.batch_size = 4`, which is a minimum for Pipeline Parallel with `pipeline_parallel_degree = 2` and `pipeline_parallel_schedule = "Interleaved1F1B"`
+- `training.data_parallel_shard_degree = 8`, resulting in global batch size 32
 - `training.steps = 3000`, `training.warmup_steps = 600`
 
-Remarks
-- This is an example series of loss-converging tests on parallelisms, not including training optimizations.
-- The default global batch size in the toml config is 64 (DP 64, local batch size 1). Given hardware resource availability, one can consider using a smaller (or larger) N. This will result in smaller (or larger, respectively) global batch size, which would possibly necessitate a smaller (or larger, respectively) learning rate (`optimizer.lr`) to keep training stability.
-
-
-| Parallelisms Dimension <br> (N = 16 by default) | Setup | Remarks |
+| Parallelism | Techniques | Remarks |
 | ----- | ----- | ----- |
-| 1D (N GPUs) | FSDP N | the 1D control set |
-| 3D (4N GPUs) | FSDP N, TP2, PP 2 | 3D test set |
-| 4D (8N GPUs) | FSDP N, TP 2, CP 2, PP 2 | 4D test set |
-| 2D (MN GPUs) <br> e.g. M=8 | FSDP N, CP M | to verify CP with a larger degree |
-
+| FSDP 8 | default | 1D control set |
+| FSDP 8, TP 2, PP 2 | torch.compile, Float8, async TP | 3D test set |
+| FSDP 8, TP 2, CP 2, PP 2 | torch.compile, Float8, async TP, Interleaved 1F1B | 4D test set |
+| FSDP 8, CP 8 | default | to verify CP with a larger degree |
 
 ### Test results
-(TBA)
+![image](../assets/images/loss_curves.png)
+
 
 [^1]: Model initialization in a sharded setting can hardly match that in a single-device setting (or a differently sharded setting), because each time a random operator is called, the underlying RNG state offset is advanced by a quantized amount, often not aligned with the amount of randomness needed, thus “wasting” different amount of randomness on differently sharded settings.
 

assets/images/llama3_loss_curves.png

@@ -1,39 +1,72 @@
-To demonstrate the effectiveness of PyTorch distributed training techniques used in torchtitan, we report both the infra metrics and loss curves of Llama 3 (8B and 70B) training on 64 A100 (80GB memory) GPUs and Llama 3.1 (405B) on 128 H100 (94GB memory).
-We report infra metrics achieved by [FSDP2](fsdp.md) (1D parallelism) under various configurations, and loss curves for both 1D parallelism (FSDP2) and 2D parallelism (FSDP2 + Tensor Parallel) training. (We only report 2D for 405B)
+We demonstrate the effectiveness of elastic distributed training using torchtitan, via experiments on Llama 3.1 8B, 70B, and 405B models, from 1D parallelism to 4D parallelism, at the scale from 8 GPUs to 512 GPUs.
 
+We ran our performance benchmarks on the [Grand Teton platform](https://engineering.fb.com/2022/10/18/open-source/ocp-summit-2022-grand-teton/), where
+- Each host has 8 NVIDIA H100 GPUs fully connected with NVLink.
+- Each H100 GPU is equipped with 96GB HBM2e with 2.4 TB/sec peak memory bandwidth.
+- Hosts are inter-connected with backend RDMA network with 400 Gb/s per GPU.
+- We used the default 500W power limit, although tuning it up to 700W TDP can potentially provide further speedups.
 
-## Llama 3.1 performance numbers
+We note that, throughout our experimentation, memory readings are stable across the whole training process[^1], whereas throughput numbers (TPS/GPU) are calculated and logged every 10 iterations, and always read at the (arbitrarily determined) 90th iteration.
 
-Below are the WPS (word per second, or more accurately, token per second) and MFU (model FLOPS utilization) results which torchtitan achieves on the 405B model released in [Llama 3.1](https://llama.meta.com/docs/model-cards-and-prompt-formats/llama3_1). The way we compute WPS and MFU can be found in `train.py`. Because the model now is larger, we run on 128 H100 GPUs to test both performance and loss curves. Below is the performance result of 405B model with optimizations we have developed. We do see OOM for 1D parallelism (FSDP2), so we only tested 2D parallelism (FSDP2 + Tensor Parallel).
+We do not report Model FLOPS Utilization (MFU) because when Float8 is enabled (on `nn.Linear` modules), both BFLOAT16 Tensor Core and FP8 Tensor Core are involved in model training, but they have different peak FLOPS and the definition of MFU under such scenario is not well-defined. We note that the 1D Llama 3.1 8B model training on 8 or 128 H100 GPUs without Float8 achieves 33% to 39% MFU[^2] (with or without torch.compile, respectively).
 
-| Model size | Batch size | Activation checkpointing | WPS | MFU | Optimizations |
-| ----- | ----- | ----- | ----- | ----- | ----- |
-| 405B | 2 | full | 109 | 29.0%[^1] | None
-| 405B | 2 | full | 177 | 23.46%[^2] | Float8
-| 405B | 2 | full | 185 | 24.52% | Float8 + Async TP
+**Table 1** 1D Parallelism (FSDP). Llama 3.1 8B model. 8 GPUs. Local batch size 2, global batch size 16. Selective activation checkpointing.
 
-Here, we use local batch size 2 (global batch size = local batch size 2 * number of FSDP ranks 16 = 32).
+| Techniques | TPS/GPU | Memory(GiB) |
+| ----- | ----: | ----: |
+| FSDP | 5,762 | 82.4 |
+| FSDP + torch.compile | 6,667 | 77.0 |
+| FSDP + torch.compile + Float8 | 8,532 | 76.8 |
 
-Next, we show the loss curves, all models are trained 3000 steps on the [C4 dataset](https://huggingface.co/datasets/allenai/c4), with global batch size 32. We have to use full AC to save memory usage. The results are shown in the picture (a TensorBoard screenshot) below.
+**Table 2** FSDP + CP + torch.compile + Float8. Llama 3.1 8B model. 8 GPUs. Local batch size 1. Full activation checkpointing.
 
-![image](../assets/images/llama3_1_405B_loss_curves.png)
+| Parallelism | Sequence Length | TPS/GPU | Memory(GiB) |
+| ----- | ----: | ----: | ----: |
+| FSDP 8, CP 1 | 32768 | 3,890 | 83.9 |
+| FSDP 4, CP 2 | 65536 | 2,540 | 84.2 |
+| FSDP 2, CP 4 | 131072 | 1,071 | 84.0 |
+| FSDP 1, CP 8 | 262144 | 548 | 84.5 |
 
-## Llama 3 performance numbers
+**Table 3** 1D Parallelism (FSDP). Llama 3.1 8B model. 128 GPUs. Local batch size 2, global batch size 256. Selective activation checkpointing.
 
-Below are the WPS and MFU results which torchtitan achieves on Llama 3 models with FSDP2 on 64 A100 (80GB) GPUs.
+| Techniques | TPS/GPU | Memory(GiB) |
+| ----- | ----: | ----: |
+| FSDP | 5,605 | 67.0 |
+| FSDP + torch.compile | 6,514 | 62.0 |
+| FSDP + torch.compile + Float8 | 8,380 | 61.8 |
 
-| Model size | Batch size | Activation checkpointing | WPS | MFU |
-| ----- | ----- | ----- | ----- | ----- |
-| 8B | 1 | selective layer | 2904 | 56.8% |
-| 8B | 1 | selective op | 2973 | 58.2% |
-| 70B | 1 | full | 331 | 51.7% |
+**Table 4** 2D parallelism (FSDP + TP) + torch.compile + Float8. Llama 3.1 70B model. 256 GPUs (FSDP 32, TP 8). Local batch size 16, global batch size 512. Full activation checkpointing.
 
-We use local batch size 1 (global batch size = local batch size 1 * number of FSDP ranks 64 = 64), because it mimics the small local batch size in large scaled training, and moreoever allows us to compare 1D (FSDP) and 2D (FSDP + TP) training under the same global batch size on both 8B and 70B Llama 3 models, without the out-of-memory (OOM) issue.
+| Techniques | TPS/GPU | Memory(GiB) |
+| ----- | ----: | ----: |
+| 2D | 829 | 71.9 |
+| 2D + AsyncTP | 876 | 67.6 |
 
-Next we show the loss curves for Llama 3 8B and Llama 3 70B training with both 1D parallelism (FSDP2) and 2D parallelism (FSDP2 + Tensor Parallel). All four models are trained the same way as mentioned above with global batch size 64. In terms of activation checkpointing (AC) configs, the Llama 3 8B training jobs use selective op AC, whereas the Llama 3 70B training jobs use full AC. The results are shown in the picture (a TensorBoard screenshot) below.
+**Table 5** 3D parallelism (FSDP + TP + PP) + torch.compile + Float8 + AsyncTP. Llama 3.1 405B model. 512 GPUs (FSDP 8, TP 8, PP8). Local batch size 32, global batch size 256. Full activation checkpointing.
 
-![image](../assets/images/llama3_loss_curves.png)
+| Schedule | TPS/GPU | Memory(GiB) |
+| ----- | ----: | ----: |
+| 1F1B | 100 | 82.5 |
+| Interleaved 1F1B | 128 | 72.7 |
 
-[^1]: We used HBM2e based lower TDP SXM H100(95GB) for our test, the actual peak TFLOPs number is between SXM and NVL, and we don't know its exact value. So this MFU number is lower than actual MFU because we use the peak number of SXM directly.
+**Table 6** 4D parallelism (FSDP + TP + PP + CP) + torch.compile + Float8 + AsyncTP + 1F1B. Llama 3.1 405B model. 512 GPUs (TP 8, PP8). Local batch size 8. Full activation checkpointing.
 
-[^2]: Since for Float8, we are not converting all the matmuls to Float8 because our fused attention implementation is not done in Float8, so this number is lower than expected.
+| Parallelism | Sequence Length | TPS/GPU | Memory(GiB) |
+| ----- | ----: | ----: | ----: |
+| FSDP 8, CP 1 | 32768 | 76 | 75.3 |
+| FSDP 4, CP 2 | 65536 | 47 | 75.9 |
+| FSDP 2, CP 4 | 131072 | 31 | 77.1 |
+| FSDP 1, CP 8 | 262144 | 16 | 84.9 |
+
+
+#### Versions used for performance testing
+| repo | commit | date |
+| --- | --- | --- |
+| torch | [1963fc8](https://github.com/pytorch/pytorch/commit/1963fc83a1c32e162162e2414f78b043f0674bae) | 2024/12/23 |
+| torchao | [eab345c](https://github.com/pytorch/ao/commit/eab345c2268a7506355d506ebfc27b5d28e5e7d0) | 2024/12/23 |
+| torchtitan | [9dec370](https://github.com/pytorch/torchtitan/commit/9dec370ad26b5f8e9a7333a0e36165018262644b) | 2024/12/26 |
+
+
+[^1]: Different PP ranks can have different peak memory usages. We take the maximum across all GPUs.
+
+[^2]: In our test we used HBM2e-based SXM H100 with lower TDP, the actual peak TFLOPs number is between SXM and NVL, and we don't know its exact value. So this MFU number is lower than actual MFU because we use the peak number of SXM directly.

assets/images/loss_curves.png

@@ -116,7 +116,11 @@ def loss_fn(pred, labels):
     model_config.vocab_size = tokenizer.n_words
     model_config.max_seq_len = job_config.training.seq_len
 
-    with FakeTensorMode() if not job_config.memory_estimation.disable_fake_mode else contextlib.nullcontext():
+    with (
+        FakeTensorMode()
+        if not job_config.memory_estimation.disable_fake_mode
+        else contextlib.nullcontext()
+    ):
 
         logger.info(
             f"Building {model_name} {job_config.model.flavor} with {model_config}"
@@ -174,8 +178,6 @@ def loss_fn(pred, labels):
                 torch.nn.utils.clip_grad_norm_(
                     model.parameters(), job_config.training.max_norm, foreach=True
                 )
-                # sync float8 amaxes and scales
-                float8_handler.sync_float8_amax_and_scale_history(model)
                 # optimizer step
                 optimizers.step()
                 lr_schedulers.step()

docs/converging.md

@@ -156,6 +156,41 @@ def __init__(
         if not self.enable_checkpoint and self.ft_manager is None:
             return
 
+<<<<<<< HEAD
+        1. even for simple PP schedules, there is a separate optimizer each PP rank.
+        rank0's optimizer would have a param_group[0] which refers to layers.0 in the original model.
+        rank1's would _also_ have a param_group[0], since it's index based, but referring to layers.1.
+        When saving, these collide and one of them is lost.  Then when reloading, only one stage can
+        restore its optimizer states, others will error.
+
+            The solution to this problem is optimizer flattening: it landed in #127071 and is enabled in TorchTitan
+            by passing the 'flatten_optimizer_state_dict' kwarg to DCP functions called in the OptimizerContainer.
+
+        2. With complex PP schedules, we have multiple model chunks per pp rank. This compounds challenge (1) by also
+        requiring us to reason about multiple 'optim' objects locally.
+
+            We solve this in the Model and Optimizer wrapper classes by flattening the state dicts from each object
+            into one state dict before saving/loading. We rely on the individual state_dicts to not collide,
+            which is gauranteed for the model by correct pipeline splitting and for the optimizer by the flattening
+            support described in (1).
+
+        3. LR schedulers also index model states like optimizers and would need to be flattened properly to support
+        resharding.  Unfortunately, the implementations of different lr_schedulers do not follow a clear pattern like
+        optimizers do, so it's hard to write a generic 'flattener' utility.
+
+            TODO: This is currently unsolved and needs a fix.
+        """
+        self.states = states
+
+        self.states.update(
+            {
+                "model": ModelWrapper(model_parts),
+                "optimizer": optimizers,
+                "dataloader": dataloader,
+                "lr_scheduler": lr_schedulers,
+            }
+        )
+=======
         self._initialize_states(
             states, dataloader, model_parts, optimizers, lr_schedulers
         )
@@ -166,6 +201,7 @@ def __init__(
         self.staging_id = None
         self.cpu_offload_state_dict = None
         self.staging_stream = torch.cuda.Stream() if self.enable_staging else None
+>>>>>>> 3430d99 ([WIP][RFC] TorchFT integration)
 
         self.folder = os.path.join(job_config.job.dump_folder, ckpt_config.folder)
         self.interval_type = (
@@ -177,7 +213,9 @@ def __init__(
         self.begin_time = 0
         self.time_sync_work = None
         self.time_sync_result = None
-        self.pg = dist.new_group(backend="gloo")
+        async_mode = ckpt_config.async_mode.lower()
+        if async_mode == AsyncMode.ASYNC or self.interval_type == IntervalType.SECONDS:
+            self.pg = dist.new_group(backend="gloo")
 
         self.keep_latest_k = ckpt_config.keep_latest_k
         self.model_weights_only = ckpt_config.model_weights_only

docs/performance.md

@@ -548,25 +548,6 @@ def __init__(self):
             action="store_true",
             help="Whether precompute float8 scales dynamically for FSDP",
         )
-        self.parser.add_argument(
-            "--float8.scaling_type_input",
-            type=str,
-            default="dynamic",
-            help="float8 scaling for input, dynamic (default) or delayed",
-            choices=["dynamic", "delayed"],
-        )
-        self.parser.add_argument(
-            "--float8.scaling_type_weight",
-            type=str,
-            default="dynamic",
-            help="float8 scaling for input, dynamic (default) or delayed",
-        )
-        self.parser.add_argument(
-            "--float8.scaling_type_grad_output",
-            type=str,
-            default="dynamic",
-            help="float8 scaling for input, dynamic (default) or delayed",
-        )
 
         # communications library settings
         self.parser.add_argument(

scripts/estimate/estimation.py

@@ -41,7 +41,7 @@ def __init__(self, job_config: JobConfig, parallel_dims: ParallelDims):
             )
             return
         try:
-            from torchao.float8 import CastConfig, Float8LinearConfig, ScalingType
+            from torchao.float8 import Float8LinearConfig
         except ImportError as e:
             raise ImportError(
                 "torchao is not installed. Please install it to use float8 linear layers."
@@ -52,14 +52,8 @@ def __init__(self, job_config: JobConfig, parallel_dims: ParallelDims):
             parallel_dims.dp_shard_enabled
             and float8_config.enable_fsdp_float8_all_gather
         )
-        scaling_type_input = ScalingType(float8_config.scaling_type_input)
-        scaling_type_weight = ScalingType(float8_config.scaling_type_weight)
-        scaling_type_grad_output = ScalingType(float8_config.scaling_type_grad_output)
         self.config = Float8LinearConfig(
             enable_fsdp_float8_all_gather=enable_fsdp_float8_all_gather,
-            cast_config_input=CastConfig(scaling_type=scaling_type_input),
-            cast_config_weight=CastConfig(scaling_type=scaling_type_weight),
-            cast_config_grad_output=CastConfig(scaling_type=scaling_type_grad_output),
         )
 
         self.enabled = True
@@ -70,15 +64,6 @@ def __init__(self, job_config: JobConfig, parallel_dims: ParallelDims):
             and float8_config.precompute_float8_dynamic_scale_for_fsdp
         )
 
-        # for sync_float8_amax_and_scale_history
-        self.delayed_scaling = (
-            scaling_type_input is ScalingType.DELAYED
-            or scaling_type_weight is ScalingType.DELAYED
-            or scaling_type_grad_output is ScalingType.DELAYED
-        )
-        self._sync_float8_amax_and_scale_history = None
-        self.compile = job_config.training.compile
-
         logger.info("Float8 training active")
 
     def convert_to_float8_training(self, model: nn.Module):
@@ -117,31 +102,3 @@ def precompute_float8_dynamic_scale_for_fsdp(
         models = [model] if isinstance(model, nn.Module) else model
         for m in models:
             precompute_float8_dynamic_scale_for_fsdp(m)
-
-    def sync_float8_amax_and_scale_history(
-        self, model: Union[nn.Module, List[nn.Module]]
-    ):
-        if not self.enabled:
-            return
-
-        if not self.delayed_scaling:
-            return
-
-        from torchao.float8 import sync_float8_amax_and_scale_history
-
-        # TODO(vkuzo): see if precalculating the modules to sync over is going to
-        # meaningfully help performance
-
-        if self._sync_float8_amax_and_scale_history is None:
-            if self.compile:
-                self._sync_float8_amax_and_scale_history = torch.compile(
-                    sync_float8_amax_and_scale_history
-                )
-            else:
-                self._sync_float8_amax_and_scale_history = (
-                    sync_float8_amax_and_scale_history
-                )
-
-        models = [model] if isinstance(model, nn.Module) else model
-        for m in models:
-            self._sync_float8_amax_and_scale_history(m)