[Distributed] Move into torchchat folder (#1203)

* Move run_dist.sh into dist_run.py

* Move distributed into torchchat folder

Author: kwen2501

dist_run.py

@@ -4,26 +4,28 @@
 # This source code is licensed under the license found in the
 # LICENSE file in the root directory of this source tree.
 
+# Example run command:
+# torchrun --nproc-per-node 4 dist_run.py llama2-7b-chat --pp 2
+# torchrun --nproc-per-node 4 dist_run.py llama3 --pp 2
+
 import argparse
 import os
 from pathlib import Path
 from types import SimpleNamespace
 from typing import Any, Dict, List, Optional, Tuple
 
-# Run command:
-# torchrun --nproc-per-node 4 dist_run.py
 import torch
 import torch.distributed as dist
 
-from distributed.logging_utils import SingletonLogger
+from torchchat.distributed.logging_utils import SingletonLogger
 
 # TODO - these are not distributed specific, consider moving to new package
-from distributed.safetensor_utils import (
+from torchchat.distributed.safetensor_utils import (
     get_hf_config_file,
     get_hf_weight_map_and_path,
     load_safetensor_weights,
 )
-from distributed.utils import (
+from torchchat.distributed.utils import (
     bytes_to_readable,
     Color as color,
     CUDATrackTime,

distributed/__init__.py

@@ -18,7 +18,7 @@
 
 from torchtune.models.llama3_2_vision._convert_weights import llama3_vision_meta_to_tune
 
-from distributed import launch_distributed, ParallelDims, parallelize_llama
+from torchchat.distributed import launch_distributed, ParallelDims, parallelize_llama
 
 from torch.distributed.device_mesh import DeviceMesh
 

run_dist.sh

@@ -0,0 +1,12 @@
+# 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.
+
+from torchchat.distributed.checkpoint import load_checkpoints_to_model
+from torchchat.distributed.logging_utils import SingletonLogger
+from torchchat.distributed.parallel_config import ParallelDims
+from torchchat.distributed.parallelize_llama import parallelize_llama
+from torchchat.distributed.utils import init_distributed
+from torchchat.distributed.world_maker import launch_distributed

torchchat/cli/builder.py

@@ -12,7 +12,7 @@
 
 import torch
 
-from distributed.logging_utils import SingletonLogger
+from torchchat.distributed.logging_utils import SingletonLogger
 logger = SingletonLogger.get_logger()
 
 

distributed/README.md renamed to torchchat/distributed/README.md

@@ -4,7 +4,7 @@
 
 from collections import defaultdict
 
-from distributed.logging_utils import SingletonLogger
+from torchchat.distributed.logging_utils import SingletonLogger
 logger = SingletonLogger.get_logger()