[fix] Example scripts miscellaneous enhancement

docs/training/multi-token-prediction.md

@@ -66,12 +66,13 @@ where:
 Here's a minimal example using the Qwen3 30B-A3B recipe with MTP enabled:
 
 ```python
-from megatron.bridge.recipes.qwen import qwen3_30b_a3b_pretrain
+from megatron.bridge.recipes.qwen.qwen3_moe import qwen3_30b_a3b_pretrain_config
 from megatron.bridge.training.pretrain import pretrain
+from megatron.bridge.training.gpt_step import forward_step
+from megatron.bridge.training.config import ConfigContainer
 
-log_dir = f"/path/to/log/dir"
+log_dir = "/path/to/log/dir"
 cfg: ConfigContainer = qwen3_30b_a3b_pretrain_config()
-cfg.logger.log_dir = log_dir
 cfg.logger.tensorboard_dir = log_dir + "/tb_logs"
 cfg.checkpoint.save = log_dir + "/checkpoints"
 cfg.checkpoint.load = log_dir + "/checkpoints"
@@ -82,10 +83,11 @@ cfg.dataset.blend=[[
 ], None]
 cfg.dataset.split="9999,8,2"
 cfg.dataset.path_to_cache = "/path/to/cache"
+# cfg.model.num_layers = 8  # train a smaller model if OOM
 # MTP Configuration
-cfg.mtp_num_layers = 1
-cfg.mtp_loss_scaling_factor = 0.1
-pretrain(cfg)
+cfg.model.mtp_num_layers = 1
+cfg.model.mtp_loss_scaling_factor = 0.1
+pretrain(cfg, forward_step)
 ```
 Follow the [DCLM Tutorial](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/tutorials/data/dclm) to prepare the training data 
 

examples/models/vlm/gemma3_vl/peft.sh

@@ -16,6 +16,10 @@
 # Workspace directory for checkpoints and results
 WORKSPACE=${WORKSPACE:-/workspace}
 
+# Before training, make sure to set WANDB_API_KEY or disable wandb logging
+# export WANDB_API_KEY=<your_wandb_api_key>
+# export WANDB_MODE=disabled
+
 # Common configurations
 PRETRAINED_CHECKPOINT=${WORKSPACE}/models/gemma-3-4b-it
 MODEL_NAME=gemma3_vl_4b

examples/models/vlm/gemma3_vl/sft.sh

@@ -16,6 +16,10 @@
 # Workspace directory for checkpoints and results
 WORKSPACE=${WORKSPACE:-/workspace}
 
+# Before training, make sure to set WANDB_API_KEY or disable wandb logging
+# export WANDB_API_KEY=<your_wandb_api_key>
+# export WANDB_MODE=disabled
+
 # Common configurations
 PRETRAINED_CHECKPOINT=${WORKSPACE}/models/gemma-3-4b-it
 MODEL_NAME=gemma3_vl_4b

examples/models/vlm/glm_45v/slurm_peft.sh

@@ -44,6 +44,10 @@
 # Workspace directory for checkpoints and results
 WORKSPACE=${WORKSPACE:-/workspace}
 
+# Before training, make sure to set WANDB_API_KEY or disable wandb logging
+# export WANDB_API_KEY=<your_wandb_api_key>
+# export WANDB_MODE=disabled
+
 # Model and training configurations
 PRETRAINED_CHECKPOINT=${WORKSPACE}/models/GLM-4.5V
 MODEL_NAME=glm_45v

examples/models/vlm/glm_45v/slurm_sft.sh

@@ -44,6 +44,10 @@
 # Workspace directory for checkpoints and results
 WORKSPACE=${WORKSPACE:-/workspace}
 
+# Before training, make sure to set WANDB_API_KEY or disable wandb logging
+# export WANDB_API_KEY=<your_wandb_api_key>
+# export WANDB_MODE=disabled
+
 # Model and training configurations
 PRETRAINED_CHECKPOINT=${WORKSPACE}/models/GLM-4.5V
 MODEL_NAME=glm_45v

examples/models/vlm/ministral3/conversion.sh

@@ -16,17 +16,20 @@
 # Workspace directory for checkpoints and results
 WORKSPACE=${WORKSPACE:-/workspace}
 
+# Note: Ministral 3 requires transformers version 5
+# pip install --upgrade transformers
+
 # Import HF → Megatron
-uv run python examples/conversion/convert_checkpoints.py import \
+python examples/conversion/convert_checkpoints.py import \
     --hf-model mistralai/Ministral-3-3B-Instruct-2512-BF16 \
     --megatron-path ${WORKSPACE}/models/Ministral-3-3B-Instruct-2512-BF16
 
 # Export Megatron → HF
-uv run python examples/conversion/convert_checkpoints.py export \
+python examples/conversion/convert_checkpoints.py export \
     --hf-model mistralai/Ministral-3-3B-Instruct-2512-BF16 \
     --megatron-path ${WORKSPACE}/models/Ministral-3-3B-Instruct-2512-BF16/iter_0000000 \
     --hf-path ${WORKSPACE}/models/Ministral-3-3B-Instruct-2512-BF16-hf-export
 
 # Round-trip validation
-uv run python -m torch.distributed.run --nproc_per_node=8 examples/conversion/hf_megatron_roundtrip_multi_gpu.py \
+python -m torch.distributed.run --nproc_per_node=8 examples/conversion/hf_megatron_roundtrip_multi_gpu.py \
     --hf-model-id mistralai/Ministral-3-3B-Instruct-2512-BF16 --tp 2 --pp 2

examples/models/vlm/ministral3/inference.sh

@@ -16,8 +16,11 @@
 # Workspace directory for checkpoints and results
 WORKSPACE=${WORKSPACE:-/workspace}
 
+# Note: Ministral 3 requires transformers version 5
+# pip install --upgrade transformers
+
 # Inference with Hugging Face checkpoints
-uv run python -m torch.distributed.run --nproc_per_node=4 examples/conversion/hf_to_megatron_generate_vlm.py \
+torchrun --nproc_per_node=4 examples/conversion/hf_to_megatron_generate_vlm.py \
     --hf_model_path mistralai/Ministral-3-3B-Instruct-2512-BF16 \
     --image_path "https://huggingface.co/nvidia/NVIDIA-Nemotron-Nano-12B-v2-VL-BF16/resolve/main/images/table.png" \
     --prompt "Describe this image." \
@@ -26,7 +29,7 @@ uv run python -m torch.distributed.run --nproc_per_node=4 examples/conversion/hf
     --pp 2
 
 # Inference with imported Megatron checkpoints
-uv run python -m torch.distributed.run --nproc_per_node=4 examples/conversion/hf_to_megatron_generate_vlm.py \
+torchrun --nproc_per_node=4 examples/conversion/hf_to_megatron_generate_vlm.py \
     --hf_model_path mistralai/Ministral-3-3B-Instruct-2512-BF16 \
     --megatron_model_path ${WORKSPACE}/models/Ministral-3-3B-Instruct-2512-BF16/iter_0000000 \
     --image_path "https://huggingface.co/nvidia/NVIDIA-Nemotron-Nano-12B-v2-VL-BF16/resolve/main/images/table.png" \
@@ -36,7 +39,7 @@ uv run python -m torch.distributed.run --nproc_per_node=4 examples/conversion/hf
     --pp 2
 
 # Inference with exported HF checkpoints
-uv run python -m torch.distributed.run --nproc_per_node=4 examples/conversion/hf_to_megatron_generate_vlm.py \
+torchrun --nproc_per_node=4 examples/conversion/hf_to_megatron_generate_vlm.py \
     --hf_model_path ${WORKSPACE}/models/Ministral-3-3B-Instruct-2512-BF16-hf-export \
     --image_path "https://huggingface.co/nvidia/NVIDIA-Nemotron-Nano-12B-v2-VL-BF16/resolve/main/images/table.png" \
     --prompt "Describe this image." \

examples/models/vlm/ministral3/peft.sh

@@ -16,6 +16,13 @@
 # Workspace directory for checkpoints and results
 WORKSPACE=${WORKSPACE:-/workspace}
 
+# Note: Ministral 3 requires transformers version 5
+# pip install --upgrade transformers
+
+# Before training, make sure to set WANDB_API_KEY or disable wandb logging
+# export WANDB_API_KEY=<your_wandb_api_key>
+# export WANDB_MODE=disabled
+
 # Common configurations
 PRETRAINED_CHECKPOINT=${WORKSPACE}/models/Ministral-3-3B-Instruct-2512-BF16
 MODEL_NAME=ministral3_3b
@@ -38,7 +45,7 @@ for config in "${PARALLELISM_CONFIGS[@]}"; do
     IFS=',' read -r TP PP <<< "$config"
 
     echo "Running LoRA finetuning with TP=$TP, PP=$PP"
-    uv run python -m torch.distributed.run --nproc_per_node=8 scripts/training/run_recipe.py \
+    torchrun --nproc_per_node=8 scripts/training/run_recipe.py \
         --recipe ${MODEL_NAME}_finetune_config \
         --step_func vlm_step \
         --peft_scheme lora \

examples/models/vlm/ministral3/sft.sh

@@ -16,6 +16,13 @@
 # Workspace directory for checkpoints and results
 WORKSPACE=${WORKSPACE:-/workspace}
 
+# Note: Ministral 3 requires transformers version 5
+# pip install --upgrade transformers
+
+# Before training, make sure to set WANDB_API_KEY or disable wandb logging
+# export WANDB_API_KEY=<your_wandb_api_key>
+# export WANDB_MODE=disabled
+
 # Common configurations
 PRETRAINED_CHECKPOINT=${WORKSPACE}/models/Ministral-3-3B-Instruct-2512-BF16
 MODEL_NAME=ministral3_3b
@@ -38,7 +45,7 @@ for config in "${PARALLELISM_CONFIGS[@]}"; do
     IFS=',' read -r TP PP <<< "$config"
 
     echo "Running full finetuning with TP=$TP, PP=$PP"
-    uv run python -m torch.distributed.run --nproc_per_node=8 scripts/training/run_recipe.py \
+    torchrun --nproc_per_node=8 scripts/training/run_recipe.py \
         --recipe ${MODEL_NAME}_finetune_config \
         --step_func vlm_step \
         checkpoint.pretrained_checkpoint=$PRETRAINED_CHECKPOINT \

examples/models/vlm/qwen3_vl/peft.sh

@@ -16,6 +16,10 @@
 # Workspace directory for checkpoints and results
 WORKSPACE=${WORKSPACE:-/workspace}
 
+# Before training, make sure to set WANDB_API_KEY or disable wandb logging
+# export WANDB_API_KEY=<your_wandb_api_key>
+# export WANDB_MODE=disabled
+
 # Common configurations for dense model finetuning
 PRETRAINED_CHECKPOINT=${WORKSPACE}/models/Qwen3-VL-8B-Instruct
 MODEL_NAME=qwen3_vl_8b

examples/models/vlm/qwen3_vl/sft.sh

@@ -16,6 +16,10 @@
 # Workspace directory for checkpoints and results
 WORKSPACE=${WORKSPACE:-/workspace}
 
+# Before training, make sure to set WANDB_API_KEY or disable wandb logging
+# export WANDB_API_KEY=<your_wandb_api_key>
+# export WANDB_MODE=disabled
+
 # Common configurations for dense model finetuning
 PRETRAINED_CHECKPOINT=${WORKSPACE}/models/Qwen3-VL-8B-Instruct
 MODEL_NAME=qwen3_vl_8b