Skip to content

Commit ae7fc2e

Browse files
committed
Add tp-size and pp-size variations to GPT-J model script
Fixes #671 Added missing tensor parallelism (tp-size) and pipeline parallelism (pp-size) variations to the get-ml-model-gptj script to match the implementation in get-ml-model-llama2. Changes: - Added tp-size.# and pp-size.# variation definitions - Set default tp-size.1 and pp-size.1 for pytorch,nvidia variation - Added MLC_NVIDIA_TP_SIZE and MLC_NVIDIA_PP_SIZE to new_env_keys This resolves the error "no scripts were found with tags: get,ml-model,gptj,_nvidia,_fp8,_tp-size.2" when users try to set MLC_NVIDIA_TP_SIZE environment variable.
1 parent cd96a8d commit ae7fc2e

File tree

1 file changed

+12
-0
lines changed

1 file changed

+12
-0
lines changed

script/get-ml-model-gptj/meta.yaml

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,8 @@ input_mapping:
1616
new_env_keys:
1717
- MLC_ML_MODEL_*
1818
- GPTJ_CHECKPOINT_PATH
19+
- MLC_NVIDIA_TP_SIZE
20+
- MLC_NVIDIA_PP_SIZE
1921
prehook_deps:
2022
- enable_if_env:
2123
MLC_TMP_REQUIRE_DOWNLOAD:
@@ -152,6 +154,8 @@ variations:
152154
pytorch,nvidia:
153155
default_variations:
154156
precision: fp8
157+
tp-size: tp-size.1
158+
pp-size: pp-size.1
155159
deps:
156160
- env:
157161
MLC_GIT_CHECKOUT_PATH_ENV_NAME: MLC_TENSORRT_LLM_CHECKOUT_PATH
@@ -253,6 +257,14 @@ variations:
253257
MLC_ML_MODEL_PRECISION: uint8
254258
MLC_ML_MODEL_WEIGHT_DATA_TYPES: uint8
255259
group: precision
260+
tp-size.#:
261+
env:
262+
MLC_NVIDIA_TP_SIZE: '#'
263+
group: tp-size
264+
pp-size.#:
265+
env:
266+
MLC_NVIDIA_PP_SIZE: '#'
267+
group: pp-size
256268
wget:
257269
add_deps_recursive:
258270
dae:

0 commit comments

Comments
 (0)