From fa57a4fd315f036e89d5754dad900d42792eafc6 Mon Sep 17 00:00:00 2001 From: caia Date: Fri, 15 May 2026 17:58:17 +0000 Subject: [PATCH 1/3] B200 fp8 perf prep + bf16 baseline sweep Adds fp8 + selective-AC config variants for Qwen3-30B-A3B and GPT-OSS-20B, docs from the bf16 batch-size sweep on 8x B200, and supporting run scripts. New configs (both wire Float8LinearConverter + Float8GroupedExpertsConverter via model_registry, compile.enable=True, mode=selective AC): - qwen3 / qwen3_30b_fp8 (local helper restructured to plumb fp8) - gpt_oss / gpt_oss_20b_fp8 Local debug: - trainer.py [GRAD-PROBE] now matches both qwen3- and gpt_oss-style expert parameter names (experts.w1/w2 and experts.mlp1_weight/mlp2_weight). Scripts: - scripts/run_{gpt_oss_20b,qwen_30b}.sbatch bf16 baselines - scripts/run_{gpt_oss_20b,qwen_30b}_fp8.sbatch fp8 variants (slurm array) - scripts/run_local.sh torchrun bypass for slurm Docs: - docs/bs_sweep_results.md baseline tps/tflops/mfu/memory across bs - docs/bs_sweep_perf_plan.md ordered optimization plan with est. tflop gain Co-Authored-By: Claude Opus 4.7 (1M context) --- docs/bs_sweep_perf_plan.md | 112 +++++++++++++++++++ docs/bs_sweep_results.md | 81 ++++++++++++++ scripts/run_gpt_oss_20b.sbatch | 58 ++++++++++ scripts/run_gpt_oss_20b_fp8.sbatch | 68 +++++++++++ scripts/run_local.sh | 89 +++++++++++++++ scripts/run_qwen_30b.sbatch | 64 +++++++++++ scripts/run_qwen_30b_fp8.sbatch | 75 +++++++++++++ torchtitan/models/gpt_oss/config_registry.py | 47 ++++++++ torchtitan/models/qwen3/config_registry.py | 72 ++++++++++++ torchtitan/trainer.py | 61 ++++++++++ 10 files changed, 727 insertions(+) create mode 100644 docs/bs_sweep_perf_plan.md create mode 100644 docs/bs_sweep_results.md create mode 100755 scripts/run_gpt_oss_20b.sbatch create mode 100644 scripts/run_gpt_oss_20b_fp8.sbatch create mode 100755 scripts/run_local.sh create mode 100755 scripts/run_qwen_30b.sbatch create mode 100644 scripts/run_qwen_30b_fp8.sbatch 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 || 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.dtype=bfloat16 \ + --training.local_batch_size="${LOCAL_BS}" diff --git a/scripts/run_local.sh b/scripts/run_local.sh new file mode 100755 index 0000000000..9830e27d47 --- /dev/null +++ b/scripts/run_local.sh @@ -0,0 +1,89 @@ +#!/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=20) +# oss-fp8 -> gpt_oss/gpt_oss_20b_fp8 (fp8 + compile + selective AC, default bs=20) +# qwen -> qwen3/qwen3_30b (bf16 + compile baseline, default bs=12) +# qwen-fp8 -> qwen3/qwen3_30b_fp8 (fp8 + compile + selective AC, default bs=12) +# +# Examples: +# ./scripts/run_local.sh oss # rerun the gpt_oss bf16 baseline at bs=20 +# ./scripts/run_local.sh qwen 12 # rerun qwen3 bf16 baseline at bs=12 +# ./scripts/run_local.sh oss-fp8 24 +# ./scripts/run_local.sh qwen-fp8 16 + +set -euo pipefail + +VARIANT="${1:?usage: $0 oss|oss-fp8|qwen|qwen-fp8 [local_batch_size]}" +LOCAL_BS="${2:-}" + +case "$VARIANT" in + oss) + MODULE=gpt_oss; CONFIG=gpt_oss_20b; DEFAULT_BS=8; TAG=oss; PREFIX=bf16-compile ;; + 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 ;; + *) + echo "unknown config '$VARIANT' — expected one of: oss, oss-fp8, qwen, qwen-fp8" >&2 + exit 2 + ;; +esac + +LOCAL_BS="${LOCAL_BS:-$DEFAULT_BS}" + +cd /data/cc/torchtitan +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 + +mkdir -p logs +TS=$(date +%Y%m%d-%H%M%S) +LOG="logs/${PREFIX}-${TAG}-b${LOCAL_BS}-${TS}.out" + +# 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__)')" +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 "" + +torchrun \ + --standalone \ + --nproc_per_node=8 \ + -m torchtitan.train \ + --module "$MODULE" \ + --config "$CONFIG" \ + --training.local_batch_size="$LOCAL_BS" \ + 2>&1 | tee "$LOG" diff --git a/scripts/run_qwen_30b.sbatch b/scripts/run_qwen_30b.sbatch new file mode 100755 index 0000000000..e3f423a7b2 --- /dev/null +++ b/scripts/run_qwen_30b.sbatch @@ -0,0 +1,64 @@ +#!/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 + +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 \ \ 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..1bfee79171 --- /dev/null +++ b/scripts/run_qwen_30b_fp8.sbatch @@ -0,0 +1,75 @@ +#!/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}" + +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.dtype=bfloat16 \ + --training.local_batch_size="${LOCAL_BS}" diff --git a/torchtitan/models/gpt_oss/config_registry.py b/torchtitan/models/gpt_oss/config_registry.py index a481daca1d..a50f7dd06e 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,48 @@ def gpt_oss_20b() -> Trainer.Config: ) +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( + expert_parallel_degree=1, + ), + 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..4b0d28ae61 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,70 @@ 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 + 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=1, + 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()], From 304bb7b940178ffa5e2341b925119f769f2a90c4 Mon Sep 17 00:00:00 2001 From: caia Date: Tue, 19 May 2026 19:33:16 +0000 Subject: [PATCH 2/3] WandB logging + fp32 variant + EP=2 for fp8 MoE configs run_local.sh: - tee full output (including early failures) to logs/--b-.out - WandB env wiring (project=mfu-b200), enabled only if wandb is importable - add oss-fp32 (full fp32, default bs=1) and deepseek (deepseek_v3_16b) variants - --tee 3 on torchrun + DTYPE passthrough config registries: - gpt_oss: add gpt_oss_20b_fp32 (selective AC, compile, bs=1) - gpt_oss_20b_fp8 + qwen3_30b_fp8: bump expert_parallel_degree 1 -> 2 (Float8GroupedExpertsConverter -> TorchAOTokenDispatcher requires EP>1 for padded token groups in quantized grouped GEMMs) sbatch tweaks: - drop redundant --training.dtype=bfloat16 from the fp8 sbatches - rename gpt_oss bf16 baseline job/log tag to "base-compile" Co-Authored-By: Claude Opus 4.7 --- scripts/run_gpt_oss_20b.sbatch | 6 +- scripts/run_gpt_oss_20b_fp8.sbatch | 1 - scripts/run_local.sh | 75 +++++++++++++++----- scripts/run_qwen_30b_fp8.sbatch | 1 - torchtitan/models/gpt_oss/config_registry.py | 36 +++++++++- torchtitan/models/qwen3/config_registry.py | 6 +- 6 files changed, 99 insertions(+), 26 deletions(-) diff --git a/scripts/run_gpt_oss_20b.sbatch b/scripts/run_gpt_oss_20b.sbatch index 4bd922fd5d..96d4e33e92 100755 --- a/scripts/run_gpt_oss_20b.sbatch +++ b/scripts/run_gpt_oss_20b.sbatch @@ -1,12 +1,12 @@ #!/bin/bash -#SBATCH --job-name=tt-gptoss-20b-bf16-compile-b8 +#SBATCH --job-name=tt-gptoss-20b-base-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-oss-b8-%j.out -#SBATCH --error=/data/cc/torchtitan/logs/bf16-compile-oss-b8-%j.err +#SBATCH --output=/data/cc/torchtitan/logs/base-compile-oss-b8-%j.out +#SBATCH --error=/data/cc/torchtitan/logs/base-compile-oss-b8-%j.err #SBATCH --partition=dev #SBATCH --qos=dev #SBATCH --account=ubuntu diff --git a/scripts/run_gpt_oss_20b_fp8.sbatch b/scripts/run_gpt_oss_20b_fp8.sbatch index f04220fb84..3e9e5985fc 100644 --- a/scripts/run_gpt_oss_20b_fp8.sbatch +++ b/scripts/run_gpt_oss_20b_fp8.sbatch @@ -64,5 +64,4 @@ srun --export=ALL torchrun \ -m torchtitan.train \ --module gpt_oss \ --config gpt_oss_20b_fp8 \ - --training.dtype=bfloat16 \ --training.local_batch_size="${LOCAL_BS}" diff --git a/scripts/run_local.sh b/scripts/run_local.sh index 9830e27d47..f087ac0ccd 100755 --- a/scripts/run_local.sh +++ b/scripts/run_local.sh @@ -6,33 +6,45 @@ # ./scripts/run_local.sh [local_batch_size] # # Configs: -# oss -> gpt_oss/gpt_oss_20b (bf16 + compile baseline, default bs=20) -# oss-fp8 -> gpt_oss/gpt_oss_20b_fp8 (fp8 + compile + selective AC, default bs=20) -# qwen -> qwen3/qwen3_30b (bf16 + compile baseline, default bs=12) -# qwen-fp8 -> qwen3/qwen3_30b_fp8 (fp8 + compile + selective AC, default bs=12) +# 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) # -# Examples: -# ./scripts/run_local.sh oss # rerun the gpt_oss bf16 baseline at bs=20 -# ./scripts/run_local.sh qwen 12 # rerun qwen3 bf16 baseline at bs=12 -# ./scripts/run_local.sh oss-fp8 24 -# ./scripts/run_local.sh qwen-fp8 16 +# 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-fp8|qwen|qwen-fp8 [local_batch_size]}" +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 ;; + 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 ;; + 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 ;; + 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 ;; + 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-fp8, qwen, qwen-fp8" >&2 + echo "unknown config '$VARIANT' — expected one of: oss, oss-fp32, oss-fp8, qwen, qwen-fp8, deepseek" >&2 exit 2 ;; esac @@ -40,6 +52,16 @@ 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}" @@ -61,9 +83,14 @@ export PYTHONFAULTHANDLER=1 export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True export TORCH_NCCL_HEARTBEAT_TIMEOUT_SEC=3600 -mkdir -p logs -TS=$(date +%Y%m%d-%H%M%S) -LOG="logs/${PREFIX}-${TAG}-b${LOCAL_BS}-${TS}.out" +# --- 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" @@ -73,17 +100,27 @@ 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" \ - 2>&1 | tee "$LOG" + --training.dtype="$DTYPE" \ + $WANDB_FLAG diff --git a/scripts/run_qwen_30b_fp8.sbatch b/scripts/run_qwen_30b_fp8.sbatch index 1bfee79171..c608181313 100644 --- a/scripts/run_qwen_30b_fp8.sbatch +++ b/scripts/run_qwen_30b_fp8.sbatch @@ -71,5 +71,4 @@ srun --export=ALL torchrun \ -m torchtitan.train \ --module qwen3 \ --config qwen3_30b_fp8 \ - --training.dtype=bfloat16 \ --training.local_batch_size="${LOCAL_BS}" diff --git a/torchtitan/models/gpt_oss/config_registry.py b/torchtitan/models/gpt_oss/config_registry.py index a50f7dd06e..b6153ba58b 100644 --- a/torchtitan/models/gpt_oss/config_registry.py +++ b/torchtitan/models/gpt_oss/config_registry.py @@ -96,6 +96,37 @@ 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"]) @@ -130,7 +161,10 @@ def gpt_oss_20b_fp8() -> Trainer.Config: steps=10000, ), parallelism=ParallelismConfig( - expert_parallel_degree=1, + # 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), diff --git a/torchtitan/models/qwen3/config_registry.py b/torchtitan/models/qwen3/config_registry.py index 4b0d28ae61..1b3b895fab 100644 --- a/torchtitan/models/qwen3/config_registry.py +++ b/torchtitan/models/qwen3/config_registry.py @@ -387,6 +387,10 @@ def _qwen3_30b_base( ) 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. @@ -416,7 +420,7 @@ def _qwen3_30b_base( parallelism=ParallelismConfig( data_parallel_shard_degree=-1, tensor_parallel_degree=1, - expert_parallel_degree=1, + expert_parallel_degree=expert_parallel_degree, pipeline_parallel_degree=1, fsdp_reshard_after_forward="default", ), From 63adf4edf7bba5bd681885451f5b38cdf547a68c Mon Sep 17 00:00:00 2001 From: caia Date: Tue, 19 May 2026 19:38:47 +0000 Subject: [PATCH 3/3] Mirror WandB wiring from run_local.sh into the 4 sbatch files run_gpt_oss_20b{,_fp8}.sbatch, run_qwen_30b{,_fp8}.sbatch: - export WANDB_PROJECT=mfu-b200, WANDB_TEAM (default caia-costello-lambd), WANDB_RUN_NAME including SLURM_JOB_ID for traceability - conditional WANDB_FLAG: only pass --metrics.enable_wandb if the wandb package is importable, so a missing install is a soft skip (mirrors the run_local.sh behavior) - flag is appended to the existing torchrun arg list, no other changes Without this, only ./scripts/run_local.sh runs logged to wandb; SLURM-launched runs silently skipped logging. Co-Authored-By: Claude Opus 4.7 --- scripts/run_gpt_oss_20b.sbatch | 16 +++++++++++++++- scripts/run_gpt_oss_20b_fp8.sbatch | 16 +++++++++++++++- scripts/run_qwen_30b.sbatch | 16 +++++++++++++++- scripts/run_qwen_30b_fp8.sbatch | 16 +++++++++++++++- 4 files changed, 60 insertions(+), 4 deletions(-) diff --git a/scripts/run_gpt_oss_20b.sbatch b/scripts/run_gpt_oss_20b.sbatch index 96d4e33e92..4dc81aa00c 100755 --- a/scripts/run_gpt_oss_20b.sbatch +++ b/scripts/run_gpt_oss_20b.sbatch @@ -37,6 +37,19 @@ 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="base-compile-oss-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 "=== GPT-OSS-20B: BF16 + compile + bs=4 (main baseline) ===" echo "Job ID: $SLURM_JOB_ID" echo "Nodes: $SLURM_NNODES" @@ -55,4 +68,5 @@ srun --export=ALL torchrun \ --module gpt_oss \ --config gpt_oss_20b \ --training.local_batch_size=8 \ - --compile.enable + --compile.enable \ + $WANDB_FLAG diff --git a/scripts/run_gpt_oss_20b_fp8.sbatch b/scripts/run_gpt_oss_20b_fp8.sbatch index 3e9e5985fc..ebdff89849 100644 --- a/scripts/run_gpt_oss_20b_fp8.sbatch +++ b/scripts/run_gpt_oss_20b_fp8.sbatch @@ -43,6 +43,19 @@ 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" @@ -64,4 +77,5 @@ srun --export=ALL torchrun \ -m torchtitan.train \ --module gpt_oss \ --config gpt_oss_20b_fp8 \ - --training.local_batch_size="${LOCAL_BS}" + --training.local_batch_size="${LOCAL_BS}" \ + $WANDB_FLAG diff --git a/scripts/run_qwen_30b.sbatch b/scripts/run_qwen_30b.sbatch index e3f423a7b2..58bc013d91 100755 --- a/scripts/run_qwen_30b.sbatch +++ b/scripts/run_qwen_30b.sbatch @@ -44,6 +44,19 @@ 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" @@ -61,4 +74,5 @@ srun --export=ALL torchrun \ -m torchtitan.train \ --module qwen3 \ --config qwen3_30b \ - --training.local_batch_size=8 \ \ No newline at end of file + --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 index c608181313..48e80ab136 100644 --- a/scripts/run_qwen_30b_fp8.sbatch +++ b/scripts/run_qwen_30b_fp8.sbatch @@ -50,6 +50,19 @@ 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" @@ -71,4 +84,5 @@ srun --export=ALL torchrun \ -m torchtitan.train \ --module qwen3 \ --config qwen3_30b_fp8 \ - --training.local_batch_size="${LOCAL_BS}" + --training.local_batch_size="${LOCAL_BS}" \ + $WANDB_FLAG