mixtral recipe#1674
Conversation
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
|
Important Review skippedAuto reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
pstjohn
force-pushed
the
pstjohn/moe-recipe
branch
from
4fa03a9 to
25b3389
Compare
pstjohn
requested review from
jstjohn,
jwilber,
savitha-eng and
trvachov
as code owners
Native PyTorch + Transformer Engine training recipe for Mixtral 8x7B that implements expert parallelism (EP) "from scratch" — 2D (dp, ep) DeviceMesh, FSDP2 fully_shard over dp, hand-written differentiable all-to-all token dispatch, explicit dense-gradient all-reduce over ep, and the fused CuteDSL MXFP8 GroupedMLP kernel — without Megatron-LM. Includes: - models/mixtral: NVMixtral TE model (attention, MoE block, token dispatcher, expert weight handling), HF<->TE conversion (convert.py/export.py), consolidated/EP-reshardable expert checkpointing (grouped_dcp.py), and tests. - recipes/mixtral_native_te: train_fsdp2_ep.py entrypoint, distributed_setup (mesh + FSDP2 wrapping + dense grad all-reduce), checkpoint save/load, THD dataloader, Hydra configs, requirements + Dockerfile (pinned CuteDSL trio), and tests (L0 sanity, EP stop/go checkpoint parity, EP=N->M reshard, 2D EP+FSDP2). - Global MoE aux-loss reduction, decoupled FusedAdam fp32 master weights, and fused save_pretrained. - Benchmarks: steady-state throughput for the four (dp, ep) layouts at MXFP8 and BF16 on 8xB300 (DCLM-baseline, max_seq=8192, token_mb=16384) and 8xB200, with a dual-axis (PFLOP/s/GPU and tokens/s/GPU) chart in the recipe README and per-run result CSVs. - perf_logger reports the true per-step peak GPU memory (max_memory_allocated) instead of the resting post-step allocation.
pstjohn
force-pushed
the
pstjohn/moe-recipe
branch
from
04a225e to
4fab9a3
Compare
…patcher Replace the moe_permute + bespoke _pad_expert_groups (+ keep-mask) path with TE's fused moe_permute_and_pad_with_probs: it groups tokens by expert, pads each group to the 256-row MXFP8 alignment (returning pad_offsets), and carries per-token routing probs so the destination-side ScaledSwiGLU applies them — the layout the fused GroupedMLP kernel wants. combine() now uses mask-type moe_unpermute with pad_offsets and merging_probs=None (weights already applied in the FFN). The empty-rank dummy-group guard is retained (empty experts stay unpadded at 0 rows). Numerically equivalent to the previous dispatcher (A/B: grouped_linear matches to ~5 sig figs, fused to ~3-4 given MXFP8 fp-reordering) and throughput-neutral at 8x1B; simplifies the expert FFN and matches TE's intended fused-MoE dispatch pattern.
Unwrap lines to the repo's 120-char ruff-format width and realign README tables via mdformat so pre-commit passes cleanly.
|
/ok to test 9b4858e |
|
One comment from an agentic review, it seems to complain about FLOPS undercounting. I am not sure how FLOPS works with EP, and unfortunately do not have time to dig at the moment, so in the name of fast feedback, here's the agent output. I did have it check itself 😅 . `Yes, after re-checking, the benchmark.py math is wrong for the metric it prints. The evidence is:
So this line is under-counting by ep_size: tokens_per_step = micro_batch * seq_len * args.dp_size For benchmark FLOPs, it should be: global_tokens_per_step = micro_batch * seq_len * dist_config.world_size One caveat: benchmark.py does not appear to write the checked-in CSVs directly, so I’d phrase the review as “benchmark.py reports wrong throughput/MFU for EP layouts,” not necessarily “the committed CSVs are |
jstjohn
left a comment
jstjohn
left a comment
There was a problem hiding this comment.
some questions about test tolerances.
Add tolerance negative controls, keep compatibility test output in tmp dirs, and centralize Mixtral recipe test path setup.
|
/ok to test d7122e7 |
Propagate the recipe root through PYTHONPATH so direct torchrun worker scripts inherit the same import path as pytest collection.
|
/ok to test 877bf85 |
3 participants
Adds
recipes/mixtral_native_te, a demonstration of how to train Mixtral 8x7b with TE's GroupedLinear and GroupedMLP layers using expert parallelism, FSDP2, and MXFP8. Benchmarks training throughput on 8xB300 under different combinations of EP and DP (FSDP2). We modify the base mixtral model to ensure support for torch distributed checkpointing, using DTensor to help coordinate the discrete, rank-local expert weights needed by TransformerEngine.