Merge branch 'main' into android-update

Author: kirklandsign

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

@@ -90,7 +90,7 @@ REQUIREMENTS_TO_INSTALL=(
 # Rely on the latest tochtune for flamingo support
 (
   set -x
-  $PIP_EXECUTABLE install git+https://github.com/pytorch/torchtune.git@18efc81dda1c537bb7c25058ff059b4623ccff58
+  $PIP_EXECUTABLE install -I git+https://github.com/pytorch/torchtune.git@d002d45e3ec700fa770d9dcc61b02c59e2507bf6
 )
 
 if [[ -x "$(command -v nvidia-smi)" ]]; then
@@ -99,3 +99,9 @@ if [[ -x "$(command -v nvidia-smi)" ]]; then
     $PYTHON_EXECUTABLE torchchat/utils/scripts/patch_triton.py
   )
 fi
+
+
+(
+  set -x
+  $PIP_EXECUTABLE install lm-eval=="0.4.2"
+)

install/install_requirements.sh

@@ -14,7 +14,6 @@ snakeviz
 sentencepiece
 numpy < 2.0
 gguf
-lm-eval==0.4.2
 blobfile
 tomli >= 1.1.0 ; python_version < "3.11"
 openai

install/requirements.txt

@@ -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()