diff --git a/docs/bs_sweep_perf_plan.md b/docs/bs_sweep_perf_plan.md new file mode 100644 index 0000000000..96c6dc68bd --- /dev/null +++ b/docs/bs_sweep_perf_plan.md @@ -0,0 +1,112 @@ +# Perf optimization plan — building on the bf16+compile baseline + +Companion to [bs_sweep_results.md](bs_sweep_results.md). This is a quick triage of perf knobs we can pull next, ordered by effort, with **rough** estimated per-GPU tflops on 8×B200 starting from those baselines. + +## Where we are today + +Steady-state, 8×B200, FSDP-only, bf16 + `torch.compile(components=["model","loss"])`, seq_len=8192: + +| Model | bs | tflops/gpu | mfu | headroom to bf16 peak (~2.25 PF/s) | headroom to fp8 peak (~4.5 PF/s) | +| --- | --- | --- | --- | --- | --- | +| gpt_oss_20b | 10 | ~573 | 25.5% | 4.0× | 7.9× | +| qwen3_30b (Qwen3-30B-A3B) | 12 | ~455 | 20.2% | 5.0× | 9.9× | + +Estimates below are **back-of-envelope only**, calibrated against the relative cost of GEMM vs comm/AC/launch in MoE LLMs at this scale. Treat the `+%` columns as "expected order of magnitude, real result needs to be measured." + +## 1. Free wins (no torchao, ~hours of work) + +### 1.1 Loosen activation checkpointing for Qwen3 + +Right now the local `qwen3_30b` helper hard-codes `mode="full"` whenever compile is on. `full` AC recomputes ~all transformer block forward in backward (~+33% compute). gpt_oss_20b already uses `selective` at bs/dev/null; then + WANDB_FLAG="--metrics.enable_wandb" +fi +# -------------------------------------------------------------------------- + +echo "=== GPT-OSS-20B: BF16 + compile + bs=4 (main baseline) ===" +echo "Job ID: $SLURM_JOB_ID" +echo "Nodes: $SLURM_NNODES" +echo "Master: $MASTER_ADDR:$MASTER_PORT" +echo "Python: $(which python)" +echo "Torch: $(python -c 'import torch;print(torch.__version__)')" +echo "" + +srun --export=ALL torchrun \ + --nproc_per_node=8 \ + --nnodes="$SLURM_NNODES" \ + --rdzv_id="$SLURM_JOB_ID" \ + --rdzv_backend=c10d \ + --rdzv_endpoint="$MASTER_ADDR:$MASTER_PORT" \ + -m torchtitan.train \ + --module gpt_oss \ + --config gpt_oss_20b \ + --training.local_batch_size=8 \ + --compile.enable \ + $WANDB_FLAG diff --git a/scripts/run_gpt_oss_20b_fp8.sbatch b/scripts/run_gpt_oss_20b_fp8.sbatch new file mode 100644 index 0000000000..ebdff89849 --- /dev/null +++ b/scripts/run_gpt_oss_20b_fp8.sbatch @@ -0,0 +1,81 @@ +#!/bin/bash +#SBATCH --job-name=tt-gptoss-20b-fp8-compile-b20 +#SBATCH --nodes=1 +#SBATCH --ntasks-per-node=1 +#SBATCH --gres=gpu:8 +#SBATCH --cpus-per-task=64 +#SBATCH --time=00:30:00 +#SBATCH --output=/data/cc/torchtitan/logs/fp8-compile-oss-b%a-%j.out +#SBATCH --error=/data/cc/torchtitan/logs/fp8-compile-oss-b%a-%j.err +#SBATCH --partition=dev +#SBATCH --qos=dev +#SBATCH --account=ubuntu + +# Sweep local_batch_size as a slurm array, sharing the same fp8 config. +# Submit with: sbatch --array=20,24,28 scripts/run_gpt_oss_20b_fp8.sbatch +#SBATCH --array=20 + +set -euo pipefail + +mkdir -p /data/cc/torchtitan/logs + +source /data/cc/torchtitan/.venv/bin/activate +cd /data/cc/torchtitan + +: "${HF_TOKEN:?HF_TOKEN must be exported before sbatch (do NOT hardcode in this file)}" +export HF_TOKEN + +export TORCHINDUCTOR_CACHE_DIR=/data/cc/torchtitan/inductor_cache +export HF_HOME=/data/cc/torchtitan/.hf_cache +export HF_DATASETS_CACHE=$HF_HOME/datasets +export HUGGINGFACE_HUB_CACHE=$HF_HOME/hub +export TRITON_CACHE_DIR=/data/cc/torchtitan/.triton_cache +export XDG_CACHE_HOME=/data/cc/torchtitan/.xdg_cache +mkdir -p "$TORCHINDUCTOR_CACHE_DIR" "$HF_HOME" "$HF_DATASETS_CACHE" "$HUGGINGFACE_HUB_CACHE" "$TRITON_CACHE_DIR" "$XDG_CACHE_HOME" + +export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1) +export MASTER_PORT=29500 +export NCCL_DEBUG=WARN +export NCCL_IB_DISABLE=0 +export PYTHONFAULTHANDLER=1 +export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True +export TORCH_NCCL_HEARTBEAT_TIMEOUT_SEC=3600 + +LOCAL_BS="${SLURM_ARRAY_TASK_ID:-20}" + +# --- WandB ---------------------------------------------------------------- +# torchtitan's WandBLogger reads these env vars (see components/metrics.py). +# Logging is enabled only if the wandb package is importable; otherwise the +# --metrics.enable_wandb flag is omitted to avoid a hard failure. +export WANDB_PROJECT="mfu-b200" +export WANDB_RUN_NAME="fp8-compile-oss-b${LOCAL_BS}-${SLURM_JOB_ID}" +export WANDB_TEAM="${WANDB_TEAM:-caia-costello-lambd}" +WANDB_FLAG="" +if python -c "import wandb" 2>/dev/null; then + WANDB_FLAG="--metrics.enable_wandb" +fi +# -------------------------------------------------------------------------- + +echo "=== GPT-OSS-20B: FP8 (rowwise linears + grouped experts) + compile + selective AC + bs=${LOCAL_BS} ===" +echo "Job ID: $SLURM_JOB_ID (array task: ${SLURM_ARRAY_TASK_ID:-n/a})" +echo "Nodes: $SLURM_NNODES" +echo "Master: $MASTER_ADDR:$MASTER_PORT" +echo "Python: $(which python)" +echo "Torch: $(python -c 'import torch;print(torch.__version__)')" +echo "TorchAO: $(python -c 'import torchao;print(torchao.__version__)' 2>/dev/null || echo 'NOT INSTALLED — fp8 will fail')" +echo "" + +# Config bakes in: compile.enable=True, components=["model","loss"], +# selective AC, Float8LinearConverter (rowwise, filter output+router.gate), +# Float8GroupedExpertsConverter. local_batch_size is overridden per array task. +srun --export=ALL torchrun \ + --nproc_per_node=8 \ + --nnodes="$SLURM_NNODES" \ + --rdzv_id="$SLURM_JOB_ID" \ + --rdzv_backend=c10d \ + --rdzv_endpoint="$MASTER_ADDR:$MASTER_PORT" \ + -m torchtitan.train \ + --module gpt_oss \ + --config gpt_oss_20b_fp8 \ + --training.local_batch_size="${LOCAL_BS}" \ + $WANDB_FLAG diff --git a/scripts/run_local.sh b/scripts/run_local.sh new file mode 100755 index 0000000000..f087ac0ccd --- /dev/null +++ b/scripts/run_local.sh @@ -0,0 +1,126 @@ +#!/bin/bash +# Plain-bash runner that bypasses slurm. Run directly on the target node +# (e.g. ml-16-b200-node-002), inside this repo's venv. +# +# Usage: +# ./scripts/run_local.sh [local_batch_size] +# +# Configs: +# oss -> gpt_oss/gpt_oss_20b (bf16 + compile baseline, default bs=8) +# oss-fp32 -> gpt_oss/gpt_oss_20b_fp32 (full fp32 + compile + selective AC, default bs=1) +# oss-fp8 -> gpt_oss/gpt_oss_20b_fp8 (fp8 + compile + selective AC, default bs=8) +# qwen -> qwen3/qwen3_30b (bf16 + compile baseline, default bs=8) +# qwen-fp8 -> qwen3/qwen3_30b_fp8 (fp8 + compile + selective AC, default bs=8) +# deepseek -> deepseek_v3/deepseek_v3_16b (bf16 + compile loss-only, EP=8, default bs=4) +# +# Logs: +# Everything (including early failures) is tee'd to logs/--b-.out +# +# WandB: +# Project is forced to "mfu-b200". To enable logging, export WANDB_API_KEY +# (or run `wandb login` once). Set WANDB_TEAM for the entity if you want +# the run to land in a team rather than your personal namespace. + +set -euo pipefail + +VARIANT="${1:?usage: $0 oss|oss-fp32|oss-fp8|qwen|qwen-fp8|deepseek [local_batch_size]}" +LOCAL_BS="${2:-}" + +# DTYPE: bfloat16 for the bf16/fp8 paths (fp8 still keeps non-quantized layers +# in bf16); float32 for the fp32 variants. +DTYPE=bfloat16 + +case "$VARIANT" in + oss) + MODULE=gpt_oss; CONFIG=gpt_oss_20b; DEFAULT_BS=8; TAG=oss; PREFIX=bf16-compile ;; + oss-fp32) + MODULE=gpt_oss; CONFIG=gpt_oss_20b_fp32; DEFAULT_BS=1; TAG=oss; PREFIX=fp32-compile; DTYPE=float32 ;; + oss-fp8) + MODULE=gpt_oss; CONFIG=gpt_oss_20b_fp8; DEFAULT_BS=8; TAG=oss; PREFIX=fp8-compile ;; + qwen) + MODULE=qwen3; CONFIG=qwen3_30b; DEFAULT_BS=8; TAG=qwen3; PREFIX=bf16-compile ;; + qwen-fp8) + MODULE=qwen3; CONFIG=qwen3_30b_fp8; DEFAULT_BS=8; TAG=qwen3; PREFIX=fp8-compile ;; + deepseek) + MODULE=deepseek_v3; CONFIG=deepseek_v3_16b; DEFAULT_BS=4; TAG=deepseek; PREFIX=bf16-compile ;; + *) + echo "unknown config '$VARIANT' — expected one of: oss, oss-fp32, oss-fp8, qwen, qwen-fp8, deepseek" >&2 + exit 2 + ;; +esac + +LOCAL_BS="${LOCAL_BS:-$DEFAULT_BS}" + +cd /data/cc/torchtitan +mkdir -p logs + +TS=$(date +%Y%m%d-%H%M%S) +LOG="logs/${PREFIX}-${TAG}-b${LOCAL_BS}-${TS}.out" + +# Redirect everything from here on to both terminal and log file, including +# stderr and any early-failure tracebacks. `exec > >(...)` swaps stdout for a +# process-substituted tee; `2>&1` folds stderr in. +exec > >(tee -a "$LOG") 2>&1 + +source /data/cc/torchtitan/.venv/bin/activate + +: "${HF_TOKEN:?HF_TOKEN must be exported before running}" +export HF_TOKEN + +export TORCHINDUCTOR_CACHE_DIR=/data/cc/torchtitan/inductor_cache +export HF_HOME=/data/cc/torchtitan/.hf_cache +export HF_DATASETS_CACHE=$HF_HOME/datasets +export HUGGINGFACE_HUB_CACHE=$HF_HOME/hub +export TRITON_CACHE_DIR=/data/cc/torchtitan/.triton_cache +export XDG_CACHE_HOME=/data/cc/torchtitan/.xdg_cache +mkdir -p "$TORCHINDUCTOR_CACHE_DIR" "$HF_HOME" "$HF_DATASETS_CACHE" "$HUGGINGFACE_HUB_CACHE" "$TRITON_CACHE_DIR" "$XDG_CACHE_HOME" + +export MASTER_ADDR=127.0.0.1 +export MASTER_PORT=29500 +export NCCL_DEBUG=WARN +export NCCL_IB_DISABLE=0 +export PYTHONFAULTHANDLER=1 +export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True +export TORCH_NCCL_HEARTBEAT_TIMEOUT_SEC=3600 + +# --- WandB ---------------------------------------------------------------- +# torchtitan's WandBLogger reads these env vars (see components/metrics.py). +# Project is auto-created on first log. +export WANDB_PROJECT="mfu-b200" +export WANDB_RUN_NAME="${PREFIX}-${TAG}-b${LOCAL_BS}-${TS}" +# Default entity. Override with `export WANDB_TEAM=...` before running. +export WANDB_TEAM="${WANDB_TEAM:-caia-costello-lambd}" +# -------------------------------------------------------------------------- + +# Only the fp8 configs need torchao; the bf16 baselines don't. +TORCHAO_NEEDED="no" +[[ "$VARIANT" == *-fp8 ]] && TORCHAO_NEEDED="yes" + +echo "=== ${MODULE}/${CONFIG} (bs=${LOCAL_BS}) ===" +echo "Host: $(hostname)" +echo "Python: $(which python)" +echo "Torch: $(python -c 'import torch;print(torch.__version__)')" +echo "WandB: project=$WANDB_PROJECT run=$WANDB_RUN_NAME $(python -c 'import wandb;print("v"+wandb.__version__)' 2>/dev/null || echo '(wandb NOT installed — set up will be skipped via --metrics.enable_wandb=false)')" +if [[ "$TORCHAO_NEEDED" == "yes" ]]; then + echo "TorchAO: $(python -c 'import torchao;print(torchao.__version__)' 2>/dev/null || echo 'NOT INSTALLED — fp8 will fail')" +fi +echo "Log: $LOG" +echo "" + +# Enable wandb only if the wandb package is importable to avoid hard failure. +# Bool flags in torchtitan's CLI are store_true: presence enables, omission disables. +WANDB_FLAG="" +if python -c "import wandb" 2>/dev/null; then + WANDB_FLAG="--metrics.enable_wandb" +fi + +torchrun \ + --standalone \ + --nproc_per_node=8 \ + --tee 3 \ + -m torchtitan.train \ + --module "$MODULE" \ + --config "$CONFIG" \ + --training.local_batch_size="$LOCAL_BS" \ + --training.dtype="$DTYPE" \ + $WANDB_FLAG diff --git a/scripts/run_qwen_30b.sbatch b/scripts/run_qwen_30b.sbatch new file mode 100755 index 0000000000..58bc013d91 --- /dev/null +++ b/scripts/run_qwen_30b.sbatch @@ -0,0 +1,78 @@ +#!/bin/bash +#SBATCH --job-name=tt-qwen3-30b-bf16-compile-b8 +#SBATCH --nodes=1 +#SBATCH --ntasks-per-node=1 +#SBATCH --gres=gpu:8 +#SBATCH --cpus-per-task=64 +#SBATCH --time=00:30:00 +#SBATCH --output=/data/cc/torchtitan/logs/bf16-compile-qwen3-b8-%j.out +#SBATCH --error=/data/cc/torchtitan/logs/bf16-compile-qwen3-b8-%j.err +#SBATCH --partition=dev +#SBATCH --qos=dev +#SBATCH --account=ubuntu + +set -euo pipefail + +mkdir -p /data/cc/torchtitan/logs + +source /data/cc/torchtitan/.venv/bin/activate +cd /data/cc/torchtitan + +: "${HF_TOKEN:?HF_TOKEN must be exported before sbatch}" +export HF_TOKEN + +export TORCHINDUCTOR_CACHE_DIR=/data/cc/torchtitan/inductor_cache +export HF_HOME=/data/cc/torchtitan/.hf_cache +export HF_DATASETS_CACHE=$HF_HOME/datasets +export HUGGINGFACE_HUB_CACHE=$HF_HOME/hub +export TRITON_CACHE_DIR=/data/cc/torchtitan/.triton_cache +export XDG_CACHE_HOME=/data/cc/torchtitan/.xdg_cache + +mkdir -p \ + "$TORCHINDUCTOR_CACHE_DIR" \ + "$HF_HOME" \ + "$HF_DATASETS_CACHE" \ + "$HUGGINGFACE_HUB_CACHE" \ + "$TRITON_CACHE_DIR" \ + "$XDG_CACHE_HOME" + +export MASTER_ADDR="$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)" +export MASTER_PORT=29500 +export NCCL_DEBUG=WARN +export NCCL_IB_DISABLE=0 +export PYTHONFAULTHANDLER=1 +export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True +export TORCH_NCCL_HEARTBEAT_TIMEOUT_SEC=3600 + +# --- WandB ---------------------------------------------------------------- +# torchtitan's WandBLogger reads these env vars (see components/metrics.py). +# Logging is enabled only if the wandb package is importable; otherwise the +# --metrics.enable_wandb flag is omitted to avoid a hard failure. +export WANDB_PROJECT="mfu-b200" +export WANDB_RUN_NAME="bf16-compile-qwen3-b8-${SLURM_JOB_ID}" +export WANDB_TEAM="${WANDB_TEAM:-caia-costello-lambd}" +WANDB_FLAG="" +if python -c "import wandb" 2>/dev/null; then + WANDB_FLAG="--metrics.enable_wandb" +fi +# -------------------------------------------------------------------------- + +echo "=== Qwen3-30B-A3B: BF16 + compile + bs=4 ===" +echo "Job ID: $SLURM_JOB_ID" +echo "Nodes: $SLURM_NNODES" +echo "Master: $MASTER_ADDR:$MASTER_PORT" +echo "Python: $(which python)" +echo "Torch: $(python -c 'import torch; print(torch.__version__)')" +echo "" + +srun --export=ALL torchrun \ + --nproc_per_node=8 \ + --nnodes="$SLURM_NNODES" \ + --rdzv_id="$SLURM_JOB_ID" \ + --rdzv_backend=c10d \ + --rdzv_endpoint="$MASTER_ADDR:$MASTER_PORT" \ + -m torchtitan.train \ + --module qwen3 \ + --config qwen3_30b \ + --training.local_batch_size=8 \ + $WANDB_FLAG \ No newline at end of file diff --git a/scripts/run_qwen_30b_fp8.sbatch b/scripts/run_qwen_30b_fp8.sbatch new file mode 100644 index 0000000000..48e80ab136 --- /dev/null +++ b/scripts/run_qwen_30b_fp8.sbatch @@ -0,0 +1,88 @@ +#!/bin/bash +#SBATCH --job-name=tt-qwen3-30b-fp8-compile +#SBATCH --nodes=1 +#SBATCH --ntasks-per-node=1 +#SBATCH --gres=gpu:8 +#SBATCH --cpus-per-task=64 +#SBATCH --time=00:30:00 +#SBATCH --output=/data/cc/torchtitan/logs/fp8-compile-qwen3-b%a-%j.out +#SBATCH --error=/data/cc/torchtitan/logs/fp8-compile-qwen3-b%a-%j.err +#SBATCH --partition=dev +#SBATCH --qos=dev +#SBATCH --account=ubuntu + +# Sweep local_batch_size as a slurm array, sharing the same fp8 config. +# Submit with: sbatch --array=12,16,20 scripts/run_qwen_30b_fp8.sbatch +#SBATCH --array=12 + +set -euo pipefail + +mkdir -p /data/cc/torchtitan/logs + +source /data/cc/torchtitan/.venv/bin/activate +cd /data/cc/torchtitan + +: "${HF_TOKEN:?HF_TOKEN must be exported before sbatch}" +export HF_TOKEN + +export TORCHINDUCTOR_CACHE_DIR=/data/cc/torchtitan/inductor_cache +export HF_HOME=/data/cc/torchtitan/.hf_cache +export HF_DATASETS_CACHE=$HF_HOME/datasets +export HUGGINGFACE_HUB_CACHE=$HF_HOME/hub +export TRITON_CACHE_DIR=/data/cc/torchtitan/.triton_cache +export XDG_CACHE_HOME=/data/cc/torchtitan/.xdg_cache + +mkdir -p \ + "$TORCHINDUCTOR_CACHE_DIR" \ + "$HF_HOME" \ + "$HF_DATASETS_CACHE" \ + "$HUGGINGFACE_HUB_CACHE" \ + "$TRITON_CACHE_DIR" \ + "$XDG_CACHE_HOME" + +export MASTER_ADDR="$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)" +export MASTER_PORT=29500 +export NCCL_DEBUG=WARN +export NCCL_IB_DISABLE=0 +export PYTHONFAULTHANDLER=1 +export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True +export TORCH_NCCL_HEARTBEAT_TIMEOUT_SEC=3600 + +LOCAL_BS="${SLURM_ARRAY_TASK_ID:-12}" + +# --- WandB ---------------------------------------------------------------- +# torchtitan's WandBLogger reads these env vars (see components/metrics.py). +# Logging is enabled only if the wandb package is importable; otherwise the +# --metrics.enable_wandb flag is omitted to avoid a hard failure. +export WANDB_PROJECT="mfu-b200" +export WANDB_RUN_NAME="fp8-compile-qwen3-b${LOCAL_BS}-${SLURM_JOB_ID}" +export WANDB_TEAM="${WANDB_TEAM:-caia-costello-lambd}" +WANDB_FLAG="" +if python -c "import wandb" 2>/dev/null; then + WANDB_FLAG="--metrics.enable_wandb" +fi +# -------------------------------------------------------------------------- + +echo "=== Qwen3-30B-A3B: FP8 (rowwise linears + grouped experts) + compile + selective AC + bs=${LOCAL_BS} ===" +echo "Job ID: $SLURM_JOB_ID (array task: ${SLURM_ARRAY_TASK_ID:-n/a})" +echo "Nodes: $SLURM_NNODES" +echo "Master: $MASTER_ADDR:$MASTER_PORT" +echo "Python: $(which python)" +echo "Torch: $(python -c 'import torch; print(torch.__version__)')" +echo "TorchAO: $(python -c 'import torchao; print(torchao.__version__)' 2>/dev/null || echo 'NOT INSTALLED — fp8 will fail')" +echo "" + +# Config bakes in: compile.enable=True, components=["model","loss"], +# selective AC, Float8LinearConverter (rowwise, filter output+router.gate), +# Float8GroupedExpertsConverter. local_batch_size is overridden per array task. +srun --export=ALL torchrun \ + --nproc_per_node=8 \ + --nnodes="$SLURM_NNODES" \ + --rdzv_id="$SLURM_JOB_ID" \ + --rdzv_backend=c10d \ + --rdzv_endpoint="$MASTER_ADDR:$MASTER_PORT" \ + -m torchtitan.train \ + --module qwen3 \ + --config qwen3_30b_fp8 \ + --training.local_batch_size="${LOCAL_BS}" \ + $WANDB_FLAG diff --git a/torchtitan/models/gpt_oss/config_registry.py b/torchtitan/models/gpt_oss/config_registry.py index a481daca1d..b6153ba58b 100644 --- a/torchtitan/models/gpt_oss/config_registry.py +++ b/torchtitan/models/gpt_oss/config_registry.py @@ -9,9 +9,14 @@ from torchtitan.components.lr_scheduler import LRSchedulersContainer from torchtitan.components.metrics import MetricsProcessor from torchtitan.components.optimizer import OptimizersContainer +from torchtitan.components.quantization import ( + Float8GroupedExpertsConverter, + Float8LinearConverter, +) from torchtitan.components.validate import Validator from torchtitan.config import ( ActivationCheckpointConfig, + CompileConfig, ParallelismConfig, TrainingConfig, ) @@ -91,6 +96,82 @@ def gpt_oss_20b() -> Trainer.Config: ) +def gpt_oss_20b_fp32() -> Trainer.Config: + """gpt_oss_20b in full fp32: params, grads, and optimizer state all fp32. + Memory-heavy (20B fp32 weights = ~80 GB even before activations); defaults + to bs=1. Selective AC + compile.""" + return Trainer.Config( + loss=ChunkedCELoss.Config(), + hf_assets_path="./assets/hf/gpt-oss-20b", + model_spec=model_registry("20b"), + dataloader=HuggingFaceTextDataLoader.Config(dataset="c4"), + optimizer=OptimizersContainer.Config(lr=8e-4), + lr_scheduler=LRSchedulersContainer.Config( + warmup_steps=2000, + decay_ratio=0.8, + decay_type="cosine", + min_lr_factor=0.1, + ), + training=TrainingConfig( + local_batch_size=1, + seq_len=8192, + steps=10000, + dtype="float32", + ), + parallelism=ParallelismConfig( + expert_parallel_degree=1, + ), + compile=CompileConfig(enable=True, components=["model", "loss"]), + checkpoint=CheckpointManager.Config(interval=500), + activation_checkpoint=ActivationCheckpointConfig(mode="selective"), + ) + + +def gpt_oss_20b_fp8() -> Trainer.Config: + """gpt_oss_20b + selective AC + rowwise fp8 linears + fp8 grouped-expert GEMMs.""" + compile_config = CompileConfig(enable=True, components=["model", "loss"]) + model_compile_enabled = ( + compile_config.enable and "model" in compile_config.components + ) + converters = [ + Float8LinearConverter.Config( + recipe_name="rowwise", + filter_fqns=["output", "router.gate"], + model_compile_enabled=model_compile_enabled, + ), + Float8GroupedExpertsConverter.Config( + model_compile_enabled=model_compile_enabled, + ), + ] + return Trainer.Config( + loss=ChunkedCELoss.Config(), + hf_assets_path="./assets/hf/gpt-oss-20b", + model_spec=model_registry("20b", converters=converters), + dataloader=HuggingFaceTextDataLoader.Config(dataset="c4"), + optimizer=OptimizersContainer.Config(lr=8e-4), + lr_scheduler=LRSchedulersContainer.Config( + warmup_steps=2000, + decay_ratio=0.8, + decay_type="cosine", + min_lr_factor=0.1, + ), + training=TrainingConfig( + local_batch_size=20, + seq_len=8192, + steps=10000, + ), + parallelism=ParallelismConfig( + # Float8GroupedExpertsConverter -> TorchAOTokenDispatcher requires + # EP>1 (padded token groups for quantized grouped GEMMs). With 8 + # GPUs and EP=2, FSDP shard fans out to the remaining 4. + expert_parallel_degree=2, + ), + compile=compile_config, + checkpoint=CheckpointManager.Config(interval=500), + activation_checkpoint=ActivationCheckpointConfig(mode="selective"), + ) + + def gpt_oss_120b() -> Trainer.Config: return Trainer.Config( loss=ChunkedCELoss.Config(), diff --git a/torchtitan/models/qwen3/config_registry.py b/torchtitan/models/qwen3/config_registry.py index bd76e3303e..1b3b895fab 100644 --- a/torchtitan/models/qwen3/config_registry.py +++ b/torchtitan/models/qwen3/config_registry.py @@ -9,8 +9,13 @@ from torchtitan.components.lr_scheduler import LRSchedulersContainer from torchtitan.components.metrics import MetricsProcessor from torchtitan.components.optimizer import OptimizersContainer, ParamGroupConfig +from torchtitan.components.quantization import ( + Float8GroupedExpertsConverter, + Float8LinearConverter, +) from torchtitan.config import ( ActivationCheckpointConfig, + CompileConfig, ParallelismConfig, TrainingConfig, ) @@ -362,3 +367,74 @@ def process_sample(sample): mode="selective", ), ) + + +# --------------------------------------------------------------------------- +# Qwen3-30B-A3B baselines (pytorch/torchtitan main) +# --------------------------------------------------------------------------- + + +def _qwen3_30b_base( + *, + fp8: bool, + local_batch_size: int, +) -> Trainer.Config: + # Compile is always on for the 30B configs we run; the CLI may still flip + # `--compile.enable=False` for ablations. + compile_config = CompileConfig(enable=True, components=["model", "loss"]) + model_compile_enabled = ( + compile_config.enable and "model" in compile_config.components + ) + + converters = None + # Float8GroupedExpertsConverter swaps in TorchAOTokenDispatcher, which + # requires EP>1 (padded token groups for quantized grouped GEMMs). When fp8 + # is on we promote EP to 2; remaining 4 GPUs fan out as FSDP shard. + expert_parallel_degree = 2 if fp8 else 1 + if fp8: + # Order matters: linear converter first so the grouped-expert converter + # sees a model config with linears already swapped. + converters = [ + Float8LinearConverter.Config( + recipe_name="rowwise", + filter_fqns=["output", "router.gate"], + model_compile_enabled=model_compile_enabled, + ), + Float8GroupedExpertsConverter.Config( + model_compile_enabled=model_compile_enabled, + ), + ] + + return Trainer.Config( + loss=ChunkedCELoss.Config(), + hf_assets_path="./assets/hf/Qwen3-30B-A3B", + model_spec=model_registry("30B-A3B", converters=converters), + dataloader=HuggingFaceTextDataLoader.Config(dataset="c4"), + optimizer=OptimizersContainer.Config(lr=3e-4), + lr_scheduler=LRSchedulersContainer.Config(warmup_steps=600), + training=TrainingConfig( + local_batch_size=local_batch_size, + seq_len=8192, + steps=3000, + ), + parallelism=ParallelismConfig( + data_parallel_shard_degree=-1, + tensor_parallel_degree=1, + expert_parallel_degree=expert_parallel_degree, + pipeline_parallel_degree=1, + fsdp_reshard_after_forward="default", + ), + compile=compile_config, + checkpoint=CheckpointManager.Config(interval=500), + activation_checkpoint=ActivationCheckpointConfig(mode="selective"), + ) + + +def qwen3_30b() -> Trainer.Config: + """bf16 + compile + selective AC. local_batch_size overridable via CLI.""" + return _qwen3_30b_base(fp8=False, local_batch_size=12) + + +def qwen3_30b_fp8() -> Trainer.Config: + """rowwise fp8 linears + fp8 grouped-expert GEMMs, on top of qwen3_30b.""" + return _qwen3_30b_base(fp8=True, local_batch_size=12) \ No newline at end of file diff --git a/torchtitan/trainer.py b/torchtitan/trainer.py index 2dcf0a36e4..7cce276cb3 100644 --- a/torchtitan/trainer.py +++ b/torchtitan/trainer.py @@ -753,6 +753,67 @@ def train_step( ) accumulated_losses.append(loss.detach()) + # DEBUG (caia): verify backward produced grads for selected params. + # Must run after all loss.backward() calls and before clip/optimizer step. + if self.step == 1 or self.step % 50 == 0: + def _local_norm(g): + if g is None: + return None + t = g.to_local() if hasattr(g, "to_local") else g + return float(t.float().norm().item()) + + try: + labels = ["attn_wq", "attn_wo", "expert_mlp1", "expert_mlp2"] + probes: dict[str, torch.nn.Parameter | None] = { + k: None for k in labels + } + + for name, p in self.model_parts[0].named_parameters(): + if "layers.0." not in name: + continue + + # Attention names vary by model: qwen3 uses qkv_linear.wq; + # gpt_oss/llama-style uses attention.wq. + if probes["attn_wq"] is None and ( + name.endswith("qkv_linear.wq.weight") + or name.endswith("attention.wq.weight") + ): + probes["attn_wq"] = p + elif probes["attn_wo"] is None and name.endswith("attention.wo.weight"): + probes["attn_wo"] = p + # Expert MLP names: qwen3/deepseek expose `experts.w1` / + # `experts.w2`; gpt_oss exposes `experts.mlp1_weight` / + # `experts.mlp2_weight` (gate_up vs down proj). + elif probes["expert_mlp1"] is None and ( + name.endswith("experts.w1") + or name.endswith("experts.mlp1_weight") + ): + probes["expert_mlp1"] = p + elif probes["expert_mlp2"] is None and ( + name.endswith("experts.w2") + or name.endswith("experts.mlp2_weight") + ): + probes["expert_mlp2"] = p + + parts = [] + for k in labels: + p = probes[k] + if p is None: + parts.append(f"{k}=None") + else: + n = _local_norm(p.grad) + parts.append(f"{k}=None" if n is None else f"{k}={n:.4f}") + + print( + f"[GRAD-PROBE step={self.step}] {' '.join(parts)}", + flush=True, + ) + except Exception as e: + print( + f"[GRAD-PROBE step={self.step}] error: {type(e).__name__}: {e}", + flush=True, + ) + with sl.log_trace_span("optim"): grad_norm = dist_utils.clip_grad_norm_( [p for m in self.model_parts for p in m.parameters()],