Skip to content

Commit b541d7a

Browse files
JRosenkranzJRosenkranz
committed
Merge branch 'main' into chunked_padding
2 parents fd2935f + 5e84730 commit b541d7a

2 files changed

Lines changed: 134 additions & 4 deletions

File tree

‎.github/workflows/run_tests.yml‎

Lines changed: 126 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,126 @@
1+
name: Non-AIU Testing
2+
3+
on: [pull_request]
4+
5+
jobs:
6+
build:
7+
runs-on: ubuntu-latest
8+
9+
steps:
10+
# Fetch full history + tags to find base version tag
11+
- uses: actions/checkout@v4
12+
with:
13+
fetch-depth: 0
14+
fetch-tags: true
15+
16+
- name: Set up Python 3.12 (with pip cache)
17+
id: setup_python
18+
uses: actions/setup-python@v5
19+
with:
20+
python-version: "3.12"
21+
cache: "pip" # enables built-in caching of ~/.cache/pip so wheels don’t re-download each run.
22+
23+
# -------------------- Restore venv cache --------------------
24+
- name: Restore Virtualenv
25+
id: cache-venv-restore
26+
uses: actions/cache/restore@v4
27+
with:
28+
path: ./.venv/
29+
key: ${{ runner.os }}-${{ steps.setup_python.outputs.python-version }}-venv-${{ hashFiles('pyproject.toml') }}
30+
restore-keys: |
31+
${{ runner.os }}-${{ steps.setup_python.outputs.python-version }}-venv-
32+
33+
# -------------------- Resolve base version from tags --------------------
34+
- name: Determine base version (from base branch tag)
35+
shell: bash
36+
run: |
37+
set -euo pipefail
38+
39+
BASE_REF="${{ github.base_ref }}"
40+
git fetch origin "${BASE_REF}" --prune --tags
41+
42+
# Try: latest v* tag that is merged into the base branch
43+
BASE_VER="$(git tag --merged "origin/${BASE_REF}" --list 'v*' \
44+
| sort -V | tail -n1 | sed 's/^v//')"
45+
46+
if [ -z "${BASE_VER}" ]; then
47+
echo "No v* tag merged into origin/${BASE_REF}; falling back to latest remote v* tag"
48+
# Fallback: latest remote v* tag (not necessarily merged)
49+
BASE_VER="$(git ls-remote --tags --sort='v:refname' origin 'v*' \
50+
| awk -F/ '{print $3}' | sed 's/\^{}//' | sed 's/^v//' | tail -n1)"
51+
fi
52+
53+
if [ -z "${BASE_VER}" ]; then
54+
BASE_VER="0.0.0"
55+
echo "Still no tag found; using ${BASE_VER}"
56+
fi
57+
58+
SHA_SHORT="${GITHUB_SHA::7}"
59+
PR_NUM="${{ github.event.number }}"
60+
# setuptools-scm will normalize `.dev` to `.dev0` — that's expected
61+
PRETEND_VER="${BASE_VER}.dev+pr${PR_NUM}.${SHA_SHORT}"
62+
63+
echo "BASE_VER=${BASE_VER}" | tee -a "$GITHUB_ENV"
64+
echo "SETUPTOOLS_SCM_PRETEND_VERSION=${PRETEND_VER}" | tee -a "$GITHUB_ENV"
65+
echo "Computed PRETEND version: ${PRETEND_VER}"
66+
67+
# -------------------- Create/Update venv & install PR code --------------------
68+
- name: Create/Update venv and install PR code (editable) with pretend version
69+
shell: bash
70+
env:
71+
SETUPTOOLS_SCM_PRETEND_VERSION: ${{ env.SETUPTOOLS_SCM_PRETEND_VERSION }}
72+
run: |
73+
set -euo pipefail
74+
if [ ! -f ".venv/bin/activate" ]; then
75+
python -m venv .venv
76+
fi
77+
. .venv/bin/activate
78+
79+
python -m pip install --upgrade pip
80+
# Editable install so your PR code is what gets tested;
81+
# setuptools-scm writes _version.py using PRETEND env
82+
pip install -e .[dev] pytest
83+
84+
echo "$VIRTUAL_ENV/bin" >> "$GITHUB_PATH"
85+
echo "VIRTUAL_ENV=$VIRTUAL_ENV" >> "$GITHUB_ENV"
86+
87+
- name: Sanity-check installed version
88+
shell: bash
89+
run: |
90+
python - <<'PY'
91+
import aiu_fms_testing_utils._version as v
92+
print("Installed __version__:", getattr(v, "__version__", None))
93+
print("version_tuple():", v.version_tuple)
94+
PY
95+
96+
# -------------------- Dataset cache --------------------
97+
- name: Cache ShareGPT dataset
98+
id: cache-sharegpt
99+
uses: actions/cache@v4
100+
with:
101+
path: ShareGPT_V3_unfiltered_cleaned_split.json
102+
key: sharegpt-${{ runner.os }}-v1
103+
104+
- name: Fetch dataset if missing
105+
if: steps.cache-sharegpt.outputs.cache-hit != 'true'
106+
run: |
107+
curl -L -o ShareGPT_V3_unfiltered_cleaned_split.json \
108+
https://huggingface.co/datasets/anon8231489123/ShareGPT_Vicuna_unfiltered/resolve/main/ShareGPT_V3_unfiltered_cleaned_split.json
109+
110+
# -------------------- Run tests --------------------
111+
- name: Test with pytest
112+
env:
113+
SHARE_GPT_DATASET_PATH: ${{ github.workspace }}/ShareGPT_V3_unfiltered_cleaned_split.json
114+
DATASET_PATH: ${{ github.workspace }}/ShareGPT_V3_unfiltered_cleaned_split.json
115+
run: |
116+
pytest -s -v -rA tests/utils
117+
pytest -s -v -rA tests/testing
118+
119+
# -------------------- Save venv cache --------------------
120+
- name: Save Virtualenv
121+
if: ${{ steps.cache-venv-restore.outputs.cache-primary-key != '' }}
122+
id: cache-venv-save
123+
uses: actions/cache/save@v4
124+
with:
125+
path: ./.venv/
126+
key: ${{ steps.cache-venv-restore.outputs.cache-primary-key }}

‎aiu_fms_testing_utils/utils/paged.py‎

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -159,10 +159,14 @@ def generate(
159159
raise ValueError("model must have a distributed_strategy")
160160

161161
kvheads = kvheads // tensor_parallel_size if kvheads > 1 else kvheads
162-
head_size = model.config.emb_dim // nheads
162+
head_size = getattr(
163+
model.config, "head_dim", model.config.emb_dim // model.config.nheads
164+
)
163165
if "fp8" in kwargs["attn_name"]:
164166
from fms_mo.aiu_addons.fp8.fp8_utils import ScaledTensor
165167

168+
already_scaled = prefill_chunk_size > 0
169+
166170
kwargs["past_key_value_states"] = [
167171
(
168172
ScaledTensor(
@@ -174,7 +178,7 @@ def generate(
174178
dtype=torch.float8_e4m3fn,
175179
),
176180
torch.tensor([1.0] * input_ids.shape[0], dtype=torch.float32),
177-
False,
181+
already_scaled,
178182
),
179183
ScaledTensor(
180184
torch.zeros(
@@ -185,7 +189,7 @@ def generate(
185189
dtype=torch.float8_e4m3fn,
186190
),
187191
torch.tensor([1.0] * input_ids.shape[0], dtype=torch.float32),
188-
False,
192+
already_scaled,
189193
),
190194
)
191195
for _ in range(model.config.nlayers)
@@ -422,7 +426,7 @@ def generate(
422426
current_kv_scales[layer_idx][0][seq_i] = t1._scale
423427
current_kv_scales[layer_idx][1][seq_i] = t2._scale
424428

425-
if seq_i != input_ids.size(0) - 1:
429+
if seq_i != input_ids.size(0) - 1 and prefill_chunk_size == 0:
426430
for layer_cache in current_kv_cache:
427431
layer_cache[0]._scaled = False
428432
layer_cache[1]._scaled = False

0 commit comments

Comments
 (0)