[integrating megablocks with open_lm] Question about megablocks + FSDP

[kernelmachine]

Hello! I'm trying to integrate megablocks with our open source LLM training library (open_lm), which uses native torch FSDP.

For some reason I am consistently seeing worse performance than my dense baselines at the same compute budget. I'm a bit stumped as to why, and I was wondering if you could provide any pointers on things to watch out for wrt integrations.

To integrate your library:

• I replace our feedforward layer with a Megablocks MoE: https://github.com/mlfoundations/open_lm/pull/115/files#diff-ae20e8018bed746c1e2ec41d171f99755ef714e018909a94b5097687ea80c3a4R230-R239

• I add a megablock_moe at every other layer (like in Switch/GShard): https://github.com/mlfoundations/open_lm/pull/115/files#diff-ae20e8018bed746c1e2ec41d171f99755ef714e018909a94b5097687ea80c3a4R319-R322

• I add the load balancing loss: https://github.com/mlfoundations/open_lm/pull/115/files#diff-52f1d1b2a425f93f57435d70bdac50ded09d6495ee2b417f999869c18cbbcbd9R198-R203

Am I missing anything else?

One hypothesis I have is that something is going wrong when I use Megablocks with FSDP.

Here is our FSDP wrapper:

https://github.com/mlfoundations/open_lm/blob/main/open_lm/main.py#L454-L462

Is there anything I need to change in the FSDP arguments to make sure FSDP doesn't interfere with all2alls? Currently it wraps the Transformer Block module, of which the MoE is a part of.

Thanks for your help!

[tgale96]

Hi! Could you share some more details on your benchmark? MoE training performance depends on a range of factors including your particular model configuration as well as how you choose to apply different kinds of sharding. FSDP could be one of the contributors here, but it'll be helpful to know what you're running to understand what might be the bottleneck.

[kernelmachine]

Sure! I am running some small experiments; 41M and 87M parameter autoregressive LMs on 8 x 40GB A100s, trained for 20B tokens a combination of the RedPajama corpus and the Pile. I am only applying moe_weight_parallelism for these models.

Actually, since I added the load balancing loss (I realized I was not including it properly originally), my validation perplexities are looking more intuitive....

41M (dense): 27.6 ppl
41M (32 experts, top-2 routing): 21.2 ppl

87M (dense): 18.7 ppl
87M (32 experts, top-2 routing): 16.6 ppl

I plan to try to sweep across a variety of settings to confirm. Is there a particular setting/benchmark that would be good for verifying performance at a small scale?

Another thing I would like to ask -- when I try to train a 160M parameter MoE (32 experts, top-2 routing) or larger with the same compute budget, I get this memory error:

File "model.py", line 240, in __init__
    self.feed_forward = MoE(moe_args)
  File "megablocks_moe.py", line 446, in __init__
    self.experts = ParallelMLP(args)
  File "megablocks_moe.py", line 118, in __init__
    self.mlp = mlp.MLP(args)
  File "megablocks/layers/mlp.py", line 110, in __init__
    self.w2.copy_(create_moe_expert_weights(
  File "megablocks/layers/mlp.py", line 36, in create_moe_expert_weights
    master_weights = torch.empty(
torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 1.12 GiB. GPU 0 has a total capacty of 39.56 GiB of which 278.56 MiB is free. Process 896914 has 4.82 GiB memory in use. Process 896913 has 4.82 GiB memory in use. Including non-PyTorch memory, this process has 4.82 GiB memory in use. Process 896911 has 4.82 GiB memory in use. Process 896912 has 4.82 GiB memory in use. Process 896907 has 6.64 GiB memory in use. Process 896908 has 3.70 GiB memory in use. Process 896909 has 4.82 GiB memory in use. Of the allocated memory 2.81 GiB is allocated by PyTorch, and 1008.85 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF

I have turned on both expert_model_parallelism and moe_weight_parallelism. Any thoughts on how to get past it?

Thanks again for your help!!

[tgale96]

You're using weight parallelism in addition to FSDP? Those are basically the same thing and I'm not sure what'd happen using both at once...

Nice results! That's starting to look pretty good!

re: memory error, that model will have <10GB of parameters total so even with no weight sharding they should fit in a single 40GB GPU. I think you're using a lot of activation memory so I'd try turning down the batch size. That being said, it looks like you're OOMing when the MoE weights are created...

[kernelmachine]

Oh! good point, I'll remove weight parallelism...

yeah, the memory error is a mystery to me, and it does OOM before it starts training...the hidden dimension of the 160M parameter model is just 768. I'll keep digging and circle back!

[kernelmachine]

Ah, the OOM was resolved by adding the device=torch.distributed.get_rank() in Arguments! Otherwise the expert weights were being initialized in rank 0 for all processes, for some reason.

[tgale96]

Ahhh that makes sense :) You probably want torch.cuda.current_devices(). I think that get_rank() will give you problems when you go to more than 8 ranks.

[kernelmachine]

Hi @tgale96 , just wanted to let you know that we've successfully integrated megablocks into openLM. Thanks a ton for your help! We can close this issue now.

For your (and others) reference, here are the results we obtain with top-2 routing, and experts at every other layer, on our A100 cluster, training on a mix of Red Pajama + the Pile:

Compute budgets

Compute type	41M	87M	160M	410M	830M
Number of nodes	1	1	1	2	4
Number of tokens	20.0B	20.0B	20.0B	20.0B	20.0B

Perplexity

Number of Experts	41M	87M	160M	410M	830M
1	27.61	18.68	14.87	10.54	9.39
8	19.85	14.66	12.26	9.82	8.84
32	20.55	15.28	14.62

Tokens/sec/GPU

Number of Experts	41M	87M	160M	410M	830M
1	141.2K	106.0K	95.5K	30.3K	16.0K
8	69.5K	66.6K	66.2K	18.5K	9.2K

Training Parameters

Number of Experts	41M	87M	160M	410M	830M
8 experts	68.9M	165.4M	360.6M	1.1B	2.4B
32 experts	164.5M	439.9M	1.0B	3.5B	7.9B

Inference Parameters

Number of Experts	41M	87M	160M	410M	830M
2 experts	45.0M	96.8M	190.7M	509.2M	1.1B

[tgale96]

Amazing! This looks excellent :)

re: throughput, would you mind sharing your MoE config for each training run? i.e., the parallelism strategy, batch/sequence length, model dims, and any other arguments you passed to MegaBlocks? Your numbers look pretty good, but I suspect we can improve quite a bit :)

[Muennighoff]

Afaict @kernelmachine integrated only the regular MoE from this repository in OpenLM but isn't Megablocks about showing that dMoE is better? Why would you not use dMoE?
cc @tgale96