Merge pull request #3996 from AI-Hypercomputer:hengtaoguo-rename

PiperOrigin-RevId: 922427130

Author: Google-ML-Automation

docs/tutorials/posttraining/knowledge_distillation.md

@@ -165,30 +165,6 @@ python3 -m tools.data_generation.generate_distillation_data_vllm \
 
 You can now fine-tune your smaller student model using supervised fine-tuning technique in MaxText.
 
-#### Fine-tune the student model using the generated dataset
-
-Example command to run fine-tuning on a TPU v6e-8:
-
-```bash
-python3 -m maxtext.trainers.post_train.sft.train_sft_deprecated \
-  run_name=${RUN_NAME?} \
-  base_output_directory=${BASE_OUTPUT_DIRECTORY?}/distillation/qwen3-32b-distill-llama3.1-8b \
-  tokenizer_path=meta-llama/Llama-3.1-8B-Instruct tokenizer_type=huggingface \
-  dataset_type=hf \
-  hf_path=parquet \
-  hf_train_files=${OUTPUT_DATASET?} \
-  train_split='train' \
-  train_data_columns=['messages'] \
-  load_parameters_path=${MAXTEXT_CKPT_PATH?}/0/items \
-  model_name=llama3.1-8b \
-  per_device_batch_size=2 \
-  steps=200 \
-  ici_expert_parallelism=-1 ici_fsdp_parallelism=4 \
-  max_target_length=2048 \
-  hf_access_token=${HF_TOKEN?} \
-  profiler=xplane
-```
-
 #### **[OPTIONAL]** Fine-tune the student model using the original dataset
 
 The checkpoint from the student model's fine-tuning (on the teacher-generated dataset) can be used for a subsequent fine-tuning stage. In this step, the student model is fine-tuned on the original dataset that was initially provided to the teacher model for generating the dataset.

docs/tutorials/posttraining/multimodal.md

@@ -130,7 +130,7 @@ Here, we use [ChartQA](https://huggingface.co/datasets/HuggingFaceM4/ChartQA) as
 export MAXTEXT_CKPT_PATH=...  # either set to an already available MaxText ckpt or to the one we just converted in the previous step
 export BASE_OUTPUT_DIRECTORY=gs://...
 export STEPS=1000
-python -m maxtext.trainers.post_train.sft.train_sft_deprecated \
+python -m maxtext.trainers.post_train.sft.train_sft_native \
     src/maxtext/configs/post_train/sft-vision-chartqa.yml \
     run_name="chartqa-sft" \
     model_name=gemma3-4b \

src/maxtext/configs/pyconfig.py

@@ -55,7 +55,7 @@
     "maxtext.trainers.post_train.dpo.train_dpo": "post_train/dpo.yml",
     "maxtext.trainers.post_train.rl.train_rl": "post_train/rl.yml",
     "maxtext.trainers.post_train.sft.train_sft": "post_train/sft.yml",
-    "maxtext.trainers.post_train.sft.train_sft_deprecated": "post_train/sft.yml",
+    "maxtext.trainers.post_train.sft.train_sft_native": "post_train/sft.yml",
     "maxtext.inference.decode": "base.yml",
     "maxtext.inference.decode_multi": "base.yml",
     "maxtext.inference.inference_microbenchmark": "base.yml",

src/maxtext/examples/multimodal_gemma3_demo.ipynb

@@ -160,7 +160,7 @@
     "STEPS=10\n",
     "PER_DEVICE_BATCH_SIZE=1\n",
     "\n",
-    "!python -m maxtext.trainers.post_train.sft.train_sft_deprecated \\\n",
+    "!python -m maxtext.trainers.post_train.sft.train_sft_native \\\n",
     "    $MAXTEXT_CONFIGS_DIR/sft-vision-chartqa.yml \\\n",
     "    run_name=$WORKLOAD_NAME \\\n",
     "    model_name=$MODEL_NAME \\\n",

…s/post_train/sft/train_sft_deprecated.py …iners/post_train/sft/train_sft_native.pysrc/maxtext/trainers/post_train/sft/train_sft_deprecated.py renamed to src/maxtext/trainers/post_train/sft/train_sft_native.py src/maxtext/trainers/post_train/sft/train_sft_deprecated.py renamed to src/maxtext/trainers/post_train/sft/train_sft_native.py

@@ -147,7 +147,7 @@ python3 -m maxtext.trainers.pre_train.train src/maxtext/configs/base.yml \
 One example command to run supervised finetuning with V3 on v5p-256. Supervised fine-tuning is only working with HuggingFace conversational datasets. And, you can customize the dataset path using the `hf_path` config and provide your access token with `hf_access_token` config.
 
 ```sh
-python3 -m maxtext.trainers.post_train.sft.train_sft_deprecated src/maxtext/configs/post_train/sft.yml \
+python3 -m maxtext.trainers.post_train.sft.train_sft_native src/maxtext/configs/post_train/sft.yml \
     base_output_directory=${BASE_OUTPUT_DIRECTORY?} \
     load_parameters_path=${SCANNED_CKPT_PATH?} \
     run_name=matmul_supervised_fine_tuning \

tests/end_to_end/tpu/deepseek/Run_DeepSeek.md

@@ -44,7 +44,7 @@ python3 -m maxtext.inference.decode "${MAXTEXT_CONFIGS_DIR:-${MAXTEXT_REPO_ROOT:
 
 # 3. SFT the MaxText converted checkpoint on ChartQA dataset
 export BASE_OUTPUT_DIRECTORY=${MODEL_BUCKET}/${MODEL_VARIATION}/unscanned/sft
-python -m maxtext.trainers.post_train.sft.train_sft_deprecated "${MAXTEXT_CONFIGS_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src/maxtext/configs}"//sft-vision-chartqa.yml \
+python -m maxtext.trainers.post_train.sft.train_sft_native "${MAXTEXT_CONFIGS_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src/maxtext/configs}"//sft-vision-chartqa.yml \
     run_name=$idx \
     model_name=$MODEL_NAME tokenizer_path="google/gemma-3-4b-pt" \
     per_device_batch_size=1 \

tests/end_to_end/tpu/gemma3/4b/test_gemma3_multimodal_sft.sh

@@ -60,7 +60,7 @@ python3 -m maxtext.trainers.pre_train.train "${MAXTEXT_CONFIGS_DIR:-${MAXTEXT_RE
 python3 -m maxtext.trainers.pre_train.train "${MAXTEXT_CONFIGS_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src/maxtext/configs}"//base.yml base_output_directory=${BASE_OUTPUT_PATH} run_name=megablox_fine_tuning model_name=${MODEL_NAME} tokenizer_type=huggingface tokenizer_path=${TOKENIZER_PATH} dataset_path=${DATASET_PATH} enable_checkpointing=true async_checkpointing=false load_parameters_path=${SCANNED_CKPT_PATH} scan_layers=True attention=flash sparse_matmul=True megablox=True dtype=bfloat16 weight_dtype=bfloat16 per_device_batch_size=4 steps=5 max_target_length=1024 ici_fsdp_parallelism=1 ici_expert_parallelism=32
 
 # Run supervised fine-tuning - megablox implementation
-python3 -m maxtext.trainers.post_train.sft.train_sft_deprecated "${MAXTEXT_CONFIGS_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src/maxtext/configs/post_train}"//sft.yml base_output_directory=${BASE_OUTPUT_PATH} run_name=megablox_supervised_fine_tuning model_name=${MODEL_NAME} tokenizer_type=huggingface tokenizer_path=${TOKENIZER_PATH} dataset_type=hf enable_checkpointing=true async_checkpointing=false load_parameters_path=${SCANNED_CKPT_PATH} scan_layers=True attention=flash sparse_matmul=True megablox=True dtype=bfloat16 weight_dtype=bfloat16 per_device_batch_size=4 steps=5 max_target_length=1024 ici_fsdp_parallelism=1 ici_expert_parallelism=32
+python3 -m maxtext.trainers.post_train.sft.train_sft_native "${MAXTEXT_CONFIGS_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src/maxtext/configs/post_train}"//sft.yml base_output_directory=${BASE_OUTPUT_PATH} run_name=megablox_supervised_fine_tuning model_name=${MODEL_NAME} tokenizer_type=huggingface tokenizer_path=${TOKENIZER_PATH} dataset_type=hf enable_checkpointing=true async_checkpointing=false load_parameters_path=${SCANNED_CKPT_PATH} scan_layers=True attention=flash sparse_matmul=True megablox=True dtype=bfloat16 weight_dtype=bfloat16 per_device_batch_size=4 steps=5 max_target_length=1024 ici_fsdp_parallelism=1 ici_expert_parallelism=32
 
 # Run decoding - megablox implementation
 # Note decode requires the access token for huggingface tokenizer even if the model is not gated

tests/end_to_end/tpu/gpt_oss/120b/test_gpt_oss.sh

@@ -66,7 +66,7 @@ python3 -m maxtext.trainers.pre_train.train "${MAXTEXT_CONFIGS_DIR:-${MAXTEXT_RE
 
 # Run supervised fine-tuning - megablox implementation
 # TODO: remove `abort_on_nan_loss=false` after b/497864549
-python3 -m maxtext.trainers.post_train.sft.train_sft_deprecated "${MAXTEXT_CONFIGS_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src/maxtext/configs/post_train}"//sft.yml base_output_directory=${BASE_OUTPUT_PATH} run_name=megablox_supervised_fine_tuning model_name=${MODEL_NAME} tokenizer_type=huggingface tokenizer_path=${TOKENIZER_PATH} dataset_type=hf enable_checkpointing=true async_checkpointing=false load_parameters_path=${SCANNED_CKPT_PATH} scan_layers=True attention=flash sparse_matmul=True megablox=True dtype=bfloat16 weight_dtype=bfloat16 per_device_batch_size=4 steps=5 max_target_length=1024 ici_fsdp_parallelism=1 ici_expert_parallelism=4 gcs_metrics=true abort_on_nan_loss=false
+python3 -m maxtext.trainers.post_train.sft.train_sft_native "${MAXTEXT_CONFIGS_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src/maxtext/configs/post_train}"//sft.yml base_output_directory=${BASE_OUTPUT_PATH} run_name=megablox_supervised_fine_tuning model_name=${MODEL_NAME} tokenizer_type=huggingface tokenizer_path=${TOKENIZER_PATH} dataset_type=hf enable_checkpointing=true async_checkpointing=false load_parameters_path=${SCANNED_CKPT_PATH} scan_layers=True attention=flash sparse_matmul=True megablox=True dtype=bfloat16 weight_dtype=bfloat16 per_device_batch_size=4 steps=5 max_target_length=1024 ici_fsdp_parallelism=1 ici_expert_parallelism=4 gcs_metrics=true abort_on_nan_loss=false
 
 # Run decoding - megablox implementation
 # Note decode requires the access token for huggingface tokenizer even if the model is not gated

tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss.sh

@@ -110,7 +110,7 @@ python3 -m maxtext.trainers.pre_train.train src/maxtext/configs/base.yml \
 One example command to run supervised finetuning with gpt-oss-20b on v5p-8. Supervised finetuning is only working with HuggingFace conversational datasets. And, you can customize the dataset path using the `hf_path` config. If using [gated dataset](https://huggingface.co/docs/hub/en/datasets-gated) or [gated model](https://huggingface.co/docs/hub/en/models-gated), you need additionally provide the access token with `hf_access_token` config.
 
 ```sh
-python3 -m maxtext.trainers.post_train.sft.train_sft_deprecated src/maxtext/configs/post_train/sft.yml \
+python3 -m maxtext.trainers.post_train.sft.train_sft_native src/maxtext/configs/post_train/sft.yml \
     base_output_directory=${BASE_OUTPUT_PATH?} \
     run_name=megablox_supervised_fine_tuning \
     model_name=gpt-oss-20b \