Merge caia-wandb-fp8-b200 into lambda-moe-tflops - #5
Open
caiacostello wants to merge 3 commits into
Open
Conversation
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) <noreply@anthropic.com>
run_local.sh:
- tee full output (including early failures) to logs/<prefix>-<tag>-b<bs>-<ts>.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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Bringing wandb + fp8 B200 changes into lambda-moe-tflops.