Experimental Design for FSDP vs EP consideration

Purpose of this experiment is to produce data that will enable us to decide between our best EP and our best FSDP config. FSDP option: d6144 · 4-of-128 experts · 48 layers — 359.6B total / 22.6B active, batch size 1024. EP option: Add LatentMoE w/ rms_norm. 2x latent compression with expert dim 6160 instead of 6160/2. Cap factor 1.3. Will run both at 4k seqlen. (8k is on-the-table for FSDP, but 4k will give a cleaner comparison here and the gap between 4k and 8k on FSDP should be minimal) ### Small scale runs We want the small scale runs to mirror the large scale dropping dynamics as cleanly as possible. Small scale runs should use EP=64 with batch size 1024 and seqlen 4k across one GB200 rack. This will keep EP, batch size, and seqlen fixed on all tests. CF should be fixed for the EP runs. Parameter norms and grads should be logged every 10 steps. The following sweep should be ran for EP 64, FSDP moe_chunk=1 (dropless), and FSDP moe_chunk=4 (minor dropping) | size | layers | active params | token budget | racks | step budget | train FLOPs | est. time | |---|---|---|---|---|---|---|---| | d768 | 8 | 54.3M | 40.7B | 1 | 9,703 | 4.81e19 | ~0.6 h | | d1024 | 12 | 145.5M | 109.1B | 1 | 26,016 | 2.39e20 | ~2.8 h | | d1536 | 16 | 434.1M | 325.6B | 1 | 77,625 | 1.58e21 | ~18.3 h | | d2048 | 22 | 1.06B | 795.5B | 4 | 47,414 | 7.81e21 | ~22.6 h | Active params exclude the embedding and lm_head. Token budget is 750 tokens per active parameter. Step budget divides the token budget by tokens/step: batch 1024 × seqlen 4096 = 4.19M tokens/step at one rack; d2048 runs 4 racks (4× batch) so 16.78M tokens/step. Layer count comes from the aspect-ratio depth rule (odd counts bumped to even). Train FLOPs are the Levanter analytic estimate (`lm_flops_per_token` × 3 for fwd+bwd × token budget), so they include the lm_head and the windowed-attention term. Est. time is the single-run wall-clock at 15% MFU on GB200 (2.5 PFLOP/s bf16-dense per GPU, 64 GPUs/rack). These runs are set at 750 tokens per active parameter, which will roughly mirror the hero run. High token count here will also counter-balance the extremely large batch size for these small models. Paloma Macro Loss eval should be on for all runs, with eval and checkpoint every 1k steps. every 3k steps for >=d1024. Evals should be ran with the model config used for training. Then, a secondary job should load each checkpoint for FSDP chunk4 and EP64 and rerun evals with dropless FSDP moe_chunk=1. This will give us a loss signal across a reasonable range of scales we can project to our target compute. ### Context Extension. Take the final trained checkpoints from d1536 and 2048. Run 10B tokens at 65k seqlen holding LR flat at 0.05 peak at batch size 64. Batch size 64*4 for largest run, to keep tokens per batch constaint. Do 4 variations: FSDP chunk 1, FSDP chunk 4, EP64, and EP64 that transitions to FSDP chunk 1 for context extension. The purpose of version 4, is to detect if EP64 will have catastrophic dropping at fewer and longer sequences per batch. ### Target Scale Comparison Run at 8 racks and target size with a LR schedule calculated from a 1T token budget, batch size 1024*8 w/ seqlen 4k. Let run elapse for first 4% over 4 hours, then kill after 40B tokens. This will capture anything unexpected with multirack QB dynamics. ### Decision Point Factors to consider: - TPS of each option at the target scale. - Projected Loss of each option at the target compute budget - Stability of EP context extension. The default conservative outcome is we stick with FSDP. We pick EP if it projects at least 5% (computed via scaling law and throughput tradeoff) lower loss at the target budget, it wins on the d2048 comparison, the context extension transition looks smooth, and there are no issues in the gradient or loss curves. This advantage would justify the added risk of higher dropping.