WIP integrate pippy's tracer frontend

Loss now runs and propagates to logger, but optimizer isn't working

ghstack-source-id: 56b0ef0
Pull Request resolved: #161

Author: wconstab

run_llama_train.sh

@@ -11,7 +11,7 @@ TRAINER_DIR=${1:-/home/$USER/local/torchtrain}
 # e.g.
 # LOG_RANK=0,1 NGPU=4 ./run_llama_train.sh
 
-NGPU=${NGPU:-"8"}
+NGPU=${NGPU:-"2"}
 
 # by default log just rank 0 output,
 LOG_RANK=${LOG_RANK:-0}

torchtrain/parallelisms/parallelize_llama.py

@@ -8,7 +8,7 @@
 from typing import Tuple
 
 import torch
-
+from pippy import annotate_split_points, Pipe, PipeSplitWrapper
 from torch.distributed._composable.fsdp import fully_shard, MixedPrecisionPolicy
 from torch.distributed._tensor import Replicate, Shard
 from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
@@ -134,7 +134,31 @@ def parallelize_llama(model, world_mesh, parallel_dims, job_config: JobConfig):
     the model must fit on GPU or CPU memory.
     """
     if parallel_dims.pp_enabled:
-        raise NotImplementedError("PP not implemented yet.")
+        pp_mesh = world_mesh["pp"]
+        stage_idx = pp_mesh.get_local_rank()
+        layers_per_rank = len(model.layers) // parallel_dims.pp
+        for i in range(1, parallel_dims.pp):
+            annotate_split_points(
+                model,
+                {
+                    f"layers.{i * layers_per_rank}": PipeSplitWrapper.SplitPoint.BEGINNING
+                },
+            )
+
+        # Get example input
+        label_shape = input_shape = (8, 2048)  # TODO
+        input_ids = torch.randint(
+            model.vocab_size, input_shape, dtype=torch.int64, device="meta"
+        )
+        labels = torch.randint(
+            model.vocab_size, label_shape, dtype=torch.int64, device="meta"
+        )
+        print("input_ids: ", input_ids.shape, input_ids.dtype)
+        print("labels: ", labels.shape, labels.dtype)
+
+        # Create a pipeline representation from the model
+        pipe = Pipe.from_tracing(model, parallel_dims.pp, example_args=(input_ids,))
+        model = pipe.get_stage_module(stage_idx)
 
     if parallel_dims.tp_enabled:
         tp_mesh = world_mesh["tp"]
@@ -233,4 +257,8 @@ def parallelize_llama(model, world_mesh, parallel_dims, job_config: JobConfig):
             logger.info(f"Applied {ac_mode} activation checkpointing to the model")
         logger.info("Applied FSDP to the model")
 
+    if parallel_dims.pp_enabled:
+        setattr(pipe.split_gm, f"submod_{stage_idx}", model)
+        return pipe
+
     return model

train.py

@@ -15,6 +15,8 @@
 
 import torch
 import torch.nn.functional as F
+from pippy.PipelineSchedule import PipelineScheduleGPipe
+from pippy.PipelineStage import PipelineStage
 from torch.distributed.elastic.multiprocessing.errors import record
 from torch.distributed.fsdp.sharded_grad_scaler import ShardedGradScaler
 from torch.distributed.tensor.parallel import loss_parallel
@@ -129,7 +131,9 @@ def main(job_config: JobConfig):
         world_size=world_size,
         enable_loss_parallel=job_config.training.enable_loss_parallel,
     )
-    torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
+    device = torch.device(f"cuda:{int(os.environ['LOCAL_RANK'])}")
+    # torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
+    torch.cuda.set_device(device)
     init_distributed(job_config)
 
     world_mesh = parallel_dims.build_mesh(device_type="cuda")
@@ -148,6 +152,14 @@ def main(job_config: JobConfig):
         dp_rank = dp_mesh.get_local_rank()
     else:
         dp_degree, dp_rank = 1, 0
+
+    if parallel_dims.pp_enabled:
+        pp_mesh = world_mesh["pp"]
+        pp_degree = pp_mesh.size()
+        pp_rank = pp_mesh.get_local_rank()
+    else:
+        pp_degree, pp_rank = 1, 0
+
     data_loader = build_dataloader_fn(
         job_config.training.dataset,
         job_config.training.dataset_path,
@@ -197,18 +209,54 @@ def main(job_config: JobConfig):
     model = models_parallelize_fns[model_name](
         model, world_mesh, parallel_dims, job_config
     )
-    # allocate sharded model on GPU and initialize weights via DTensor
-    model.to_empty(device="cuda")
-    model.init_weights()
-
-    # build optimizer after applying parallelisms to the model
-    optimizer = build_optimizer(model, job_config)
-    scheduler = get_lr_scheduler(optimizer, job_config)
+    if parallel_dims.pp_enabled:
+        pipe_meta = model
+        model = pipe_meta.get_stage_module(pp_rank)
 
     # build grad scaler which is effective only when mixed precision training
     # is enabled with fp16 param dtype under FSDP
     scaler = build_grad_scaler(model)
 
+    def loss_fn(pred, labels):
+        with (
+            loss_parallel()
+            if parallel_dims.loss_parallel_enabled
+            else contextlib.nullcontext()
+        ):
+            loss = F.cross_entropy(pred.flatten(0, 1), labels.flatten(0, 1))
+
+            # backward on scaled loss to create scaled gradients
+            scaler.scale(loss)
+        return loss
+
+    # TODO(whc) everything below needs to become a function that can be applied to each 'virtual stage' of PP, if
+    # there are virtual stages
+    if parallel_dims.pp_enabled:
+        stage = PipelineStage(
+            pipe=pipe_meta,
+            stage_index=pp_rank,
+            device=device,
+            group=pp_mesh.get_group(),
+        )
+        pp_schedule = PipelineScheduleGPipe(
+            stage,
+            n_microbatches=parallel_dims.pp,
+            loss_fn=loss_fn,
+        )
+        model.to_empty(device="cuda")
+    else:
+        # if PP is enabled, we can't use init_weights. instead, we have to rely on offline creating an initial checkpoint
+        # and loading it to get initialization values.  This is becuase the init_weights functions are written assuming
+        # the whole model (all its weights, or FQNs) exist on one rank.  In PP, the init_weights on stage1 might crash
+        # becuase it can't find "embedding" layer, for example.
+
+        # allocate sharded model on GPU and initialize weights via DTensor
+        model.to_empty(device="cuda")
+        model.init_weights()
+
+    # build optimizer after applying parallelisms to the model
+    optimizer = build_optimizer(model, job_config)
+    scheduler = get_lr_scheduler(optimizer, job_config)
     metric_logger = build_metric_logger(job_config)
 
     # torch.compile model for improved performance
@@ -278,21 +326,32 @@ def main(job_config: JobConfig):
 
             input_ids = input_ids.cuda()
             labels = labels.cuda()
-
             optimizer.zero_grad()
 
-            # forward
-            pred = model(input_ids)
+            if parallel_dims.pp_enabled:
+                # pipeline F/Loss/B
+                is_last_stage = pp_mesh.get_local_rank() == pp_mesh.size() - 1
 
-            with (
-                loss_parallel()
-                if parallel_dims.loss_parallel_enabled
-                else contextlib.nullcontext()
-            ):
-                loss = F.cross_entropy(pred.flatten(0, 1), labels.flatten(0, 1))
+                if pp_mesh.get_local_rank() == 0:
+                    pp_schedule.step(input_ids)
+                elif is_last_stage:
+                    losses = []
+                    pp_schedule.step(target=labels, losses=losses)
+                else:
+                    schedule.step()
+
+                # accumulate losses across pipeline microbatches
+                current_loss = (
+                    torch.mean(torch.stack(losses)).item() if is_last_stage else -1.0
+                )
+            else:
+                # non-pipeline F/Loss/B
+                pred = model(input_ids)
+
+                loss = loss_fn(pred, labels)
+                loss.backward()
 
-                # backward on scaled loss to create scaled gradients
-                scaler.scale(loss).backward()
+                current_loss = loss.item()
 
             # clip gradients (after unscaling gradients of the optimizer's params)
             scaler.unscale_(optimizer)
@@ -309,7 +368,6 @@ def main(job_config: JobConfig):
             # updates the scale for next iteration
             scaler.update()
 
-            current_loss = loss.item()
             losses_since_last_log.append(current_loss)
 
             # log metrics

train_configs/debug_model.toml

@@ -38,7 +38,7 @@ max_norm = 1.0  # grad norm clipping
 steps = 10
 data_parallel_degree = -1
 tensor_parallel_degree = 1
-pipeline_parallel_degree = 1
+pipeline_parallel_degree = 2
 fp8_linear = ""
 compile = false
 dataset = "alpaca"   # supported datasets = alpaca (52K), minipile (1M), c4 (177M)