Skip to content

Commit 1a3af81

Browse files
Jack-KhuuJack-Khuu
authored
Merge branch 'main' into nlp/cache-start-pos
2 parents 70ab45b + b2d8f2a commit 1a3af81

File tree

11 files changed

+215
-30
lines changed

11 files changed

+215
-30
lines changed

‎.github/workflows/run-readme-pr-linuxaarch64.yml

Lines changed: 124 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,124 @@
1+
name: Run the README instructions - with stories - on Linux aarch64
2+
3+
on:
4+
pull_request:
5+
push:
6+
branches:
7+
- main
8+
workflow_dispatch:
9+
10+
jobs:
11+
test-readme-cpu:
12+
uses: pytorch/test-infra/.github/workflows/linux_job.yml@main
13+
with:
14+
runner: linux-aarch64
15+
gpu-arch-type: cuda
16+
gpu-arch-version: "12.1"
17+
timeout: 60
18+
script: |
19+
echo "::group::Print machine info"
20+
uname -a
21+
echo "::endgroup::"
22+
23+
echo "::group::Install newer objcopy that supports --set-section-alignment"
24+
yum install -y devtoolset-10-binutils
25+
export PATH=/opt/rh/devtoolset-10/root/usr/bin/:$PATH
26+
echo "::endgroup::"
27+
28+
TORCHCHAT_DEVICE=cpu .ci/scripts/run-docs readme
29+
30+
echo "::group::Completion"
31+
echo "tests complete"
32+
echo "*******************************************"
33+
echo "::endgroup::"
34+
35+
test-quantization-cpu:
36+
uses: pytorch/test-infra/.github/workflows/linux_job.yml@main
37+
with:
38+
runner: linux-aarch64
39+
gpu-arch-type: cuda
40+
gpu-arch-version: "12.1"
41+
timeout: 60
42+
script: |
43+
echo "::group::Print machine info"
44+
uname -a
45+
echo "::endgroup::"
46+
47+
echo "::group::Install newer objcopy that supports --set-section-alignment"
48+
yum install -y devtoolset-10-binutils
49+
export PATH=/opt/rh/devtoolset-10/root/usr/bin/:$PATH
50+
echo "::endgroup::"
51+
52+
TORCHCHAT_DEVICE=cpu .ci/scripts/run-docs quantization
53+
54+
test-gguf-cpu:
55+
uses: pytorch/test-infra/.github/workflows/linux_job.yml@main
56+
with:
57+
runner: linux-aarch64
58+
gpu-arch-type: cuda
59+
gpu-arch-version: "12.1"
60+
timeout: 60
61+
script: |
62+
echo "::group::Print machine info"
63+
uname -a
64+
echo "::endgroup::"
65+
66+
echo "::group::Install newer objcopy that supports --set-section-alignment"
67+
yum install -y devtoolset-10-binutils
68+
export PATH=/opt/rh/devtoolset-10/root/usr/bin/:$PATH
69+
echo "::endgroup::"
70+
71+
TORCHCHAT_DEVICE=cpu .ci/scripts/run-docs gguf
72+
73+
echo "::group::Completion"
74+
echo "tests complete"
75+
echo "*******************************************"
76+
echo "::endgroup::"
77+
78+
test-advanced-cpu:
79+
uses: pytorch/test-infra/.github/workflows/linux_job.yml@main
80+
with:
81+
runner: linux-aarch64
82+
gpu-arch-type: cuda
83+
gpu-arch-version: "12.1"
84+
timeout: 60
85+
script: |
86+
echo "::group::Print machine info"
87+
uname -a
88+
echo "::endgroup::"
89+
90+
echo "::group::Install newer objcopy that supports --set-section-alignment"
91+
yum install -y devtoolset-10-binutils
92+
export PATH=/opt/rh/devtoolset-10/root/usr/bin/:$PATH
93+
echo "::endgroup::"
94+
95+
TORCHCHAT_DEVICE=cpu .ci/scripts/run-docs advanced
96+
97+
echo "::group::Completion"
98+
echo "tests complete"
99+
echo "*******************************************"
100+
echo "::endgroup::"
101+
102+
test-evaluation-cpu:
103+
uses: pytorch/test-infra/.github/workflows/linux_job.yml@main
104+
with:
105+
runner: linux-aarch64
106+
gpu-arch-type: cuda
107+
gpu-arch-version: "12.1"
108+
timeout: 60
109+
script: |
110+
echo "::group::Print machine info"
111+
uname -a
112+
echo "::endgroup::"
113+
114+
echo "::group::Install newer objcopy that supports --set-section-alignment"
115+
yum install -y devtoolset-10-binutils
116+
export PATH=/opt/rh/devtoolset-10/root/usr/bin/:$PATH
117+
echo "::endgroup::"
118+
119+
TORCHCHAT_DEVICE=cpu .ci/scripts/run-docs evaluation
120+
121+
echo "::group::Completion"
122+
echo "tests complete"
123+
echo "*******************************************"
124+
echo "::endgroup::"

‎.github/workflows/run-readme-pr-mps.yml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@ jobs:
1010
uses: pytorch/test-infra/.github/workflows/macos_job.yml@main
1111
with:
1212
runner: macos-m1-14
13-
timeout: 50
13+
timeout: 60
1414
script: |
1515
conda create -y -n test-readme-mps-macos python=3.10.11 llvm-openmp
1616
conda activate test-readme-mps-macos

‎docs/ADVANCED-USERS.md

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -479,7 +479,7 @@ in a Python-free environment with AOT Inductor and ExecuTorch.
479479
| Hardware | OS | Eager | Eager + Compile | AOT Compile | ET Runtime |
480480
|-----|------|-----|-----|-----|-----|
481481
| x86 | Linux | ✅ | ✅ | ✅ | ✅ |
482-
| aarch64 | Linux | n/t | n/t | n/t | n/t |
482+
| aarch64 | Linux | ✅ | ✅ | ✅ | n/t |
483483
| aarch64 | macOS | ✅ | ✅ | ✅ | ✅ |
484484
| AMD GPU | Linux | ✅ | ✅ | ✅ | ❌|
485485
| Nvidia GPU | Linux | ✅ | ✅ | ✅ | ❌|
@@ -490,7 +490,7 @@ in a Python-free environment with AOT Inductor and ExecuTorch.
490490
| Mobile GPU (Vulkan) | Android | ❌|❌|❌| ✅ |
491491
| CoreML | iOS | ❌|❌|❌| ✅ |
492492
| Hexagon DSP | Android | ❌|❌|❌| ✅ |
493-
| Raspberry Pi 4/5 | Raspbian | n/t | n/t | n/t | ✅ |
493+
| Raspberry Pi 4/5 | Raspbian | ✅ | ✅ | ✅ | ✅ |
494494
| Raspberry Pi 4/5 | Android | ❌ | ❌ | ❌ | n/t |
495495
| ARM 32b (up to v7) | any | ❌|❌|❌|❌|
496496

‎install/install_requirements.sh

Lines changed: 26 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -51,19 +51,13 @@ echo "Using pip executable: $PIP_EXECUTABLE"
5151
# NOTE: If a newly-fetched version of the executorch repo changes the value of
5252
# PYTORCH_NIGHTLY_VERSION, you should re-run this script to install the necessary
5353
# package versions.
54-
PYTORCH_NIGHTLY_VERSION=dev20241218
54+
PYTORCH_NIGHTLY_VERSION=dev20250119
5555

5656
# Nightly version for torchvision
57-
VISION_NIGHTLY_VERSION=dev20241218
57+
VISION_NIGHTLY_VERSION=dev20250119
5858

5959
# Nightly version for torchtune
60-
TUNE_NIGHTLY_VERSION=dev20241218
61-
62-
# Uninstall triton, as nightly will depend on pytorch-triton, which is one and the same
63-
(
64-
set -x
65-
$PIP_EXECUTABLE uninstall -y triton
66-
)
60+
TUNE_NIGHTLY_VERSION=dev20250119
6761

6862
# The pip repository that hosts nightly torch packages. cpu by default.
6963
# If cuda is available, based on presence of nvidia-smi, install the pytorch nightly
@@ -74,16 +68,28 @@ then
7468
elif [[ -x "$(command -v rocminfo)" ]];
7569
then
7670
TORCH_NIGHTLY_URL="https://download.pytorch.org/whl/nightly/rocm6.2"
71+
elif [[ -x "$(command -v xpu-smi)" ]];
72+
then
73+
TORCH_NIGHTLY_URL="https://download.pytorch.org/whl/nightly/xpu"
7774
else
7875
TORCH_NIGHTLY_URL="https://download.pytorch.org/whl/nightly/cpu"
7976
fi
8077

8178
# pip packages needed by exir.
82-
REQUIREMENTS_TO_INSTALL=(
83-
torch=="2.6.0.${PYTORCH_NIGHTLY_VERSION}"
84-
torchvision=="0.22.0.${VISION_NIGHTLY_VERSION}"
85-
torchtune=="0.5.0.${TUNE_NIGHTLY_VERSION}"
86-
)
79+
if [[ -x "$(command -v xpu-smi)" ]];
80+
then
81+
REQUIREMENTS_TO_INSTALL=(
82+
torch=="2.7.0.${PYTORCH_NIGHTLY_VERSION}"
83+
torchvision=="0.22.0.${VISION_NIGHTLY_VERSION}"
84+
torchtune=="0.6.0"
85+
)
86+
else
87+
REQUIREMENTS_TO_INSTALL=(
88+
torch=="2.7.0.${PYTORCH_NIGHTLY_VERSION}"
89+
torchvision=="0.22.0.${VISION_NIGHTLY_VERSION}"
90+
torchtune=="0.6.0.${TUNE_NIGHTLY_VERSION}"
91+
)
92+
fi
8793

8894
#
8995
# First install requirements in install/requirements.txt. Older torch may be
@@ -95,6 +101,12 @@ REQUIREMENTS_TO_INSTALL=(
95101
$PIP_EXECUTABLE install -r install/requirements.txt --extra-index-url "${TORCH_NIGHTLY_URL}"
96102
)
97103

104+
# Uninstall triton, as nightly will depend on pytorch-triton, which is one and the same
105+
(
106+
set -x
107+
$PIP_EXECUTABLE uninstall -y triton
108+
)
109+
98110
# Install the requirements. --extra-index-url tells pip to look for package
99111
# versions on the provided URL if they aren't available on the default URL.
100112
(
@@ -116,8 +128,6 @@ if [[ -x "$(command -v nvidia-smi)" ]]; then
116128
$PYTHON_EXECUTABLE torchchat/utils/scripts/patch_triton.py
117129
)
118130
fi
119-
120-
121131
(
122132
set -x
123133
$PIP_EXECUTABLE install evaluate=="0.4.3" lm-eval=="0.4.2" psutil=="6.0.0"

‎torchchat/cli/builder.py

Lines changed: 19 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -69,10 +69,16 @@ class BuilderArgs:
6969
prefill_possible: bool = False
7070
dynamic_shapes: bool = False
7171
max_seq_length: Optional[int] = None
72+
attention_backend: str = "math"
7273

7374
def __post_init__(self):
7475
if self.device is None:
75-
self.device = "cuda" if torch.cuda.is_available() else "cpu"
76+
if torch.cuda.is_available():
77+
self.device = "cuda"
78+
elif torch.xpu.is_available():
79+
self.device = "xpu"
80+
else:
81+
self.device = "cpu"
7682

7783
if not (
7884
(self.checkpoint_path and self.checkpoint_path.is_file())
@@ -178,6 +184,17 @@ def from_args(cls, args: argparse.Namespace) -> "BuilderArgs":
178184
pp = getattr(args, "pp", 1)
179185
tp = getattr(args, "tp", 1)
180186
chpt_from = getattr(args, "chpt_from", "hf")
187+
sdp_backend_dict = {
188+
'math': torch.nn.attention.SDPBackend.MATH,
189+
'flash_attention': torch.nn.attention.SDPBackend.FLASH_ATTENTION,
190+
'efficient_attention': torch.nn.attention.SDPBackend.EFFICIENT_ATTENTION,
191+
'cudnn_attention': torch.nn.attention.SDPBackend.CUDNN_ATTENTION,
192+
}
193+
attention_backend = sdp_backend_dict[args.attention_backend]
194+
if args.device == "cpu" and (args.attention_backend == "efficient_attention"
195+
or args.attention_backend == "cudnn_attention"):
196+
print(f"Warning: {args.attention_backend} is not supported on CPU. Using math instead.")
197+
attention_backend = torch.nn.attention.SDPBackend.MATH
181198
return cls(
182199
checkpoint_dir=checkpoint_dir,
183200
checkpoint_path=checkpoint_path,
@@ -202,6 +219,7 @@ def from_args(cls, args: argparse.Namespace) -> "BuilderArgs":
202219
is_chat_model=is_chat_model,
203220
dynamic_shapes=getattr(args, "dynamic_shapes", False),
204221
max_seq_length=getattr(args, "max_seq_length", None),
222+
attention_backend=attention_backend,
205223
)
206224

207225
@classmethod

‎torchchat/cli/cli.py

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -176,8 +176,15 @@ def _add_model_config_args(parser, verb: str) -> None:
176176
"--device",
177177
type=str,
178178
default=None,
179-
choices=["fast", "cpu", "cuda", "mps"],
180-
help="Hardware device to use. Options: fast, cpu, cuda, mps",
179+
choices=["fast", "cpu", "cuda", "mps", "xpu"],
180+
help="Hardware device to use. Options: fast, cpu, cuda, mps, xpu",
181+
)
182+
model_config_parser.add_argument(
183+
"--attention-backend",
184+
type=str,
185+
default="math",
186+
choices=["math", "flash_attention", "efficient_attention", "cudnn_attention"],
187+
help="SDPBackend to use. Options: MATH, FLASH_ATTENTION, EFFICIENT_ATTENTION, CUDNN_ATTENTION",
181188
)
182189

183190

‎torchchat/generate.py

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626
import torch.distributed as dist
2727
import torch.multiprocessing as mp
2828
from torch.distributed.pipelining import PipelineStage, ScheduleGPipe
29+
from torch._C import _SDPBackend as SDPBackend
2930

3031
from PIL import Image
3132

@@ -531,6 +532,7 @@ def decode_n_tokens(
531532
callback=lambda _: _,
532533
eos_token_id: int = 2,
533534
eot_id: Optional[int] = None,
535+
attention_backend: SDPBackend = torch.nn.attention.SDPBackend.MATH,
534536
**sampling_kwargs,
535537
):
536538
new_tokens, new_probs = [], []
@@ -539,7 +541,7 @@ def decode_n_tokens(
539541
num_new_tokens - 1
540542
): # -1 to save space to run an EoS if dont generate it naturally
541543
# Actually better for Inductor to codegen attention here
542-
with torch.nn.attention.sdpa_kernel([torch.nn.attention.SDPBackend.MATH]):
544+
with torch.nn.attention.sdpa_kernel([attention_backend]):
543545

544546
out_token = cur_token.clone()
545547
next_token, next_prob = self.decode_one_token(
@@ -683,6 +685,7 @@ def generate(
683685
sequential_prefill=True,
684686
callback=lambda x: x,
685687
max_seq_length: int,
688+
attention_backend: str = "math",
686689
seed: Optional[int] = None,
687690
**sampling_kwargs,
688691
) -> torch.Tensor:
@@ -799,6 +802,7 @@ def generate(
799802
if self.is_llama3_model
800803
else None
801804
),
805+
attention_backend=attention_backend,
802806
**sampling_kwargs,
803807
):
804808
generated_tokens.append(generated_token.view(-1))
@@ -1186,6 +1190,7 @@ def callback(x, *, done_generating=False):
11861190
start_pos=start_pos,
11871191
skip_cache_setup=not is_first_sample,
11881192
max_seq_length=max_seq_length,
1193+
attention_backend=self.builder_args.attention_backend,
11891194
)
11901195
if generator_args.chat_mode:
11911196
start_pos += encoded.size(0)
@@ -1205,8 +1210,10 @@ def callback(x, *, done_generating=False):
12051210
if hasattr(prof, "export_chrome_trace"):
12061211
if self.builder_args.device == "cpu":
12071212
print(prof.key_averages().table(sort_by="self_cpu_time_total"))
1208-
else:
1213+
elif self.builder_args.device == "cuda":
12091214
print(prof.key_averages().table(sort_by="self_cuda_time_total"))
1215+
else:
1216+
print(prof.key_averages().table(sort_by="self_xpu_time_total"))
12101217
prof.export_chrome_trace(f"{self.profile}.json")
12111218

12121219
if start_pos >= max_seq_length:
@@ -1291,6 +1298,9 @@ def callback(x, *, done_generating=False):
12911298
)
12921299
if torch.cuda.is_available():
12931300
print(f"Memory used: {torch.cuda.max_memory_reserved() / 1e9:.02f} GB")
1301+
if torch.xpu.is_available():
1302+
print(f"Memory used: {torch.xpu.max_memory_reserved() / 1e9:.02f} GB")
1303+
12941304

12951305

12961306
class DistributedGenerator(LocalGenerator):
@@ -1617,6 +1627,8 @@ def run_generator(
16171627
)
16181628
if torch.cuda.is_available():
16191629
torch.cuda.reset_peak_memory_stats()
1630+
if torch.xpu.is_available():
1631+
torch.xpu.reset_peak_memory_stats()
16201632

16211633
for _ in gen.chat(generator_args):
16221634
pass

‎torchchat/model.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1025,7 +1025,7 @@ def apply_rotary_emb(x: Tensor, freqs_cis: Tensor) -> Tensor:
10251025
# For quantized_decomposed ops
10261026
from executorch.kernels import quantized # no-qa
10271027
# For llama::sdpa_with_kv_cache.out, preprocess ops
1028-
from executorch.extension.llm.custom_ops import sdpa_with_kv_cache # no-qa
1028+
from executorch.extension.llm.custom_ops import custom_ops # no-qa
10291029

10301030
class PTEModel(nn.Module):
10311031
def __init__(self, config, path) -> None:
@@ -1062,5 +1062,6 @@ def forward(self, x, input_pos):
10621062
def setup_caches(self, max_batch_size, max_seq_length):
10631063
pass
10641064

1065-
except:
1065+
except Exception as e:
1066+
print(f"Warning: PTEModel (ExecuTorch) not available with exception: {e}")
10661067
pass

0 commit comments

Comments
 (0)