Skip to content

Commit 9cb7785

Browse files
Yijia Jinjetstream authors
authored andcommitted
Refactor: Execute Maxtext inference as a module via python3 -m
Updated the recommended execution command for the inference script from `python3 MaxText/decode.py` to `python3 -m MaxText.decode`. This change utilizes Python's standard module execution mechanism (`-m`), ensuring the script runs within the proper package context (`MaxText`). This improves consistency and robustness of import resolution. No changes were made to the source code itself. PiperOrigin-RevId: 749207719
1 parent 51ab94c commit 9cb7785

File tree

3 files changed

+3
-3
lines changed

3 files changed

+3
-3
lines changed

benchmarks/mlperf/scripts/launch_microbenchmark.sh

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@ export run_dir=${base_output_dir}/microbenchmark/${run_name}/${experiment_time}/
1212
echo "run_dir: ${run_dir}"
1313
gsutil cp ${config_file_path} ${run_dir}/
1414

15-
python3 MaxText/inference_microbenchmark.py \
15+
python3 -m MaxText.inference_microbenchmark \
1616
${config_file_path} \
1717
model_name=${model_name} \
1818
tokenizer_path=assets/tokenizer.llama2 \

benchmarks/mlperf/scripts/launch_server.sh

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ source ./${config}.sh
33
export run_name=${model_name}_${tpu}_${attention}_ici_${ici_tensor_parallelism}-${ici_autoregressive_parallelism}_${reshape_q}_${quant_mode}_pbs${per_device_batch_size}_${compute_axis_order//,/}-${prefill_cache_axis_order//,/}-${ar_cache_axis_order//,/}
44

55
cd /maxtext
6-
python3 MaxText/maxengine_server.py \
6+
python3 -m MaxText.maxengine_server \
77
${config_file_path} \
88
model_name=${model_name} \
99
tokenizer_path=assets/tokenizer.llama2 \

jetstream/tools/maxtext/model_ckpt_finetune_with_aqt.sh

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -50,7 +50,7 @@ export CONVERTED_CHECKPOINT=${MODEL_BUCKET}/${MODEL}/${MODEL_VARIATION}/${idx}/0
5050
# Fine tune the converted model checkpoints with AQT.
5151
export RUN_NAME=finetune_aqt_${idx}
5252

53-
python3 MaxText/train.py \
53+
python3 -m MaxText.train \
5454
MaxText/configs/base.yml \
5555
run_name=${RUN_NAME} \
5656
base_output_directory=${BASE_OUTPUT_DIRECTORY} \

0 commit comments

Comments
 (0)