|
| 1 | +# Copyright (c) Meta Platforms, Inc. and affiliates. |
| 2 | +# All rights reserved. |
| 3 | +# |
| 4 | +# This source code is licensed under the BSD-style license found in the |
| 5 | +# LICENSE file in the root directory of this source tree. |
| 6 | + |
| 7 | +import time |
| 8 | + |
| 9 | +import torch |
| 10 | + |
| 11 | +from autoparallel.api import AutoParallel |
| 12 | + |
| 13 | +from torch.distributed import DeviceMesh |
| 14 | +from torch.distributed.tensor.placement_types import Replicate, Shard |
| 15 | + |
| 16 | +from torchtitan.config_manager import JobConfig |
| 17 | +from torchtitan.distributed import ParallelDims |
| 18 | + |
| 19 | +from torchtitan.tools.logging import logger |
| 20 | + |
| 21 | + |
| 22 | +def parallelize_llama( |
| 23 | + model, |
| 24 | + world_mesh: DeviceMesh, |
| 25 | + parallel_dims: ParallelDims, |
| 26 | + job_config: JobConfig, |
| 27 | +): |
| 28 | + """ |
| 29 | + Apply tensor parallelism, activation checkpointing, torch.compile, and data |
| 30 | + parallelism to the model. |
| 31 | +
|
| 32 | + NOTE: The passed-in model preferably should be on meta device. Otherwise, |
| 33 | + the model must fit on GPU or CPU memory. |
| 34 | + """ |
| 35 | + |
| 36 | + def input_fn(): |
| 37 | + global_batch_size = job_config.training.global_batch_size |
| 38 | + if global_batch_size < 0: |
| 39 | + # This global batch size results in 1 gradient accumulation |
| 40 | + # step. |
| 41 | + dp_degree = world_mesh["dp"].size() |
| 42 | + global_batch_size = job_config.training.local_batch_size * dp_degree |
| 43 | + return torch.rand( |
| 44 | + (global_batch_size, job_config.training.seq_len), device="cuda" |
| 45 | + ) |
| 46 | + |
| 47 | + # TODO make autop work correctly with different combinations of DP, DP+TP, TP, and support DDP / HSDP |
| 48 | + assert ( |
| 49 | + len(world_mesh.shape) == 2 |
| 50 | + ), "Only support 2D mesh (DP, TP) for now- OK if one has size=1" |
| 51 | + assert parallel_dims.dp_shard_enabled is True, "DDP not supported yet" |
| 52 | + assert parallel_dims.dp_replicate_enabled is False, "DDP not supported yet" |
| 53 | + assert parallel_dims.cp_enabled is False, "CP not supported yet" |
| 54 | + assert parallel_dims.pp_enabled is False, "PP not supported yet" |
| 55 | + |
| 56 | + # bail out |
| 57 | + # model = model_fn() |
| 58 | + # return model |
| 59 | + |
| 60 | + autop = AutoParallel(model, input_fn, world_mesh) |
| 61 | + autop.add_parameter_memory_constraint(low=None, high=None) |
| 62 | + |
| 63 | + x_sharding = (Shard(0), Replicate()) |
| 64 | + |
| 65 | + autop.add_input_constraints([x_sharding]) |
| 66 | + autop.add_output_constraints([x_sharding]) |
| 67 | + t0 = time.time() |
| 68 | + sharding_placement = autop.optimize_placement() |
| 69 | + t1 = time.time() |
| 70 | + logger.info(f"AutoParallel took {t1 - t0} seconds") |
| 71 | + parallel_mod = autop.apply_placement(sharding_placement) |
| 72 | + |
| 73 | + if job_config.training.compile: |
| 74 | + torch._inductor.config.reorder_for_peak_memory = False |
| 75 | + parallel_mod = torch.compile(parallel_mod, fullgraph=True) |
| 76 | + |
| 77 | + return parallel_mod |
0 commit comments